D9 重链剖分+线段树 • EH's Blog

题目情境#

题目链接：https://www.luogu.com.cn/problem/P3384 ↗

P3384 【模板】重链剖分/树链剖分

题目描述

如题，已知一棵包含 $N$ 个结点的树（连通且无环），每个节点上包含一个数值，需要支持以下操作：

1 x y z，表示将树从 $x$ 到 $y$ 结点最短路径上所有节点的值都加上 $z$ 。
2 x y，表示求树从 $x$ 到 $y$ 结点最短路径上所有节点的值之和。
3 x z，表示将以 $x$ 为根节点的子树内所有节点值都加上 $z$ 。
4 x，表示求以 $x$ 为根节点的子树内所有节点值之和。

输入格式

第一行包含 $4$ 个正整数 $N,M,R,P$ ，分别表示树的结点个数、操作个数、根节点序号和取模数（即所有的输出结果均对此取模）。

接下来一行包含 $N$ 个非负整数，分别依次表示各个节点上初始的数值。

接下来 $N-1$ 行每行包含两个整数 $x,y$ ，表示点 $x$ 和点 $y$ 之间连有一条边（保证无环且连通）。

接下来 $M$ 行每行包含若干个正整数，每行表示一个操作。

输出格式

输出包含若干行，分别依次表示每个操作 $2$ 或操作 $4$ 所得的结果（对 $P$ 取模）。

输入 #1

plaintext

输出 #1

2
21

plaintext

【数据规模】

对于 $100\%$ 的数据： $1\le N \leq {10}^5$ ， $1\le M \leq {10}^5$ ， $1\le R\le N$ ， $1\le P \le 2^{30}$ 。所有输入的数均在 int 范围内。

【样例说明】

树的结构如下：

各个操作如下：

故输出应依次为 $2$ 和 $21$ 。

树链剖分部分和上面一样，多了一个“同链打标记”，可以使用线段树维护。总时间复杂度： ${O(mlog^2n)}$

算法解析：#

重链剖分的核心思想是把树拍平成一个序列，用 dfs 序把树上的路径/子树操作变成线段树上的区间操作。

第一遍 DFS（dfs1），自底向上求出每个点的：

$fa[u]$ ：父节点
$dep[u]$ ：深度
$siz[u]$ ：子树大小
$son[u]$ ：重儿子，即子树最大的那个儿子（若有多个取其一）

第二遍 DFS（dfs2），按”先走重儿子”的顺序给每个点分配 dfs 序 $dfn[u]$ ，并记录 $top[u]$ （所在重链的链顶）。因为每次优先进入重儿子，所以同一条重链上的点在 $dfn$ 中是连续的一段区间——这正是”一条链 = 一段区间”的关键。

剖分完成后：

子树操作：以 $x$ 为根的子树，其所有点在 $dfn$ 中恰好是连续区间 $[dfn[x],\ dfn[x]+siz[x]-1]$ ，直接对线段树做一次区间加/区间和即可， $O(\log n)$ 。
路径操作（ $x\to y$ $x \to y$ ）：不断把 $x,y$ $x, y$ 中链顶更深的那个”整条链”一起处理：
1. 若 $top[x]\neq top[y]$ ，比较 $dep[top[x]]$ 与 $dep[top[y]]$ ，取更深的一方（假设是 $x$ ），对区间 $[dfn[top[x]], dfn[x]]$ 做线段树区间操作，然后 $x=fa[top[x]]$ （跳到链顶的父节点）；
2. 重复直到 $x,y$ 位于同一条重链；
3. 最后对 $[dfn[x],dfn[y]]$ （深度小的在前）做一次区间操作。

因为每次跳链都会让链顶深度至少减半，最多跳 $O(\log n)$ 条链，每次线段树操作 $O(\log n)$ ，所以单次路径操作复杂度 $O(\log^2 n)$ 。线段树只需支持区间加与区间求和即可（本题是纯加法，不需要处理懒标记下传顺序；加乘混合时才要特殊处理）。

py代码实现#

略

C++代码实现#

// 树链剖分 O(mlognlogn)
#include<bits/stdc++.h>
using namespace std;

#define int long long
const int N=100010;
int n,m,root,P,w[N],a,b,c,t;
vector<int> e[N];
int fa[N],dep[N],siz[N],son[N];
int top[N],id[N],nw[N],cnt; //树链

#define lc u<<1
#define rc u<<1|1
struct tree{
  int l,r; 
  int sum,add;
}tr[N*4]; //线段树

void dfs1(int u,int f){ //搞fa,dep,siz,son
  fa[u]=f,dep[u]=dep[f]+1,siz[u]=1;
  for(int v:e[u]){
    if(v==f) continue;
    dfs1(v,u);
    siz[u]+=siz[v];
    if(siz[son[u]]<siz[v]) son[u]=v; 
  }
}
void dfs2(int u,int tp){ //搞top,id,nw
  top[u]=tp,id[u]=++cnt,nw[cnt]=w[u];
  if(son[u]) dfs2(son[u],tp);
  for(int v:e[u]){
    if(v==fa[u]||v==son[u])continue;
    dfs2(v,v);
  }
}

void pushup(int u){
  tr[u].sum=tr[lc].sum+tr[rc].sum;
}
void pushdown(int u){
  if(tr[u].add){
    tr[lc].sum+=tr[u].add*(tr[lc].r-tr[lc].l+1);
    tr[rc].sum+=tr[u].add*(tr[rc].r-tr[rc].l+1);
    tr[lc].add+=tr[u].add;
    tr[rc].add+=tr[u].add;
    tr[u].add=0;
  }
}
void build(int u,int l,int r){ //构建线段树
  tr[u]={l,r,nw[l],0};
  if(l==r) return;
  int mid=l+r>>1;
  build(lc,l,mid);
  build(rc,mid+1,r);
  pushup(u);
}
void change(int u,int x,int y,int k){ //线段树修改
  if(x>tr[u].r||y<tr[u].l) return;
  if(x<=tr[u].l&&tr[u].r<=y){
    tr[u].sum+=k*(tr[u].r-tr[u].l+1);
    tr[u].add+=k;
    return;
  }
  pushdown(u);
  change(lc,x,y,k);
  change(rc,x,y,k);
  pushup(u);
}
void change_path(int u,int v,int k){ //修改路径
  while(top[u]!=top[v]){
    if(dep[top[u]]<dep[top[v]]) swap(u,v);
    change(1,id[top[u]],id[u],k);
    u=fa[top[u]];
  }
  if(dep[u]<dep[v]) swap(u,v);
  change(1,id[v],id[u],k); //最后一段
}
void change_tree(int u,int k){ //修改子树
  change(1,id[u],id[u]+siz[u]-1,k);
}
int query(int u,int x,int y){ //线段树查询
  if(x>tr[u].r||y<tr[u].l) return 0;
  if(x<=tr[u].l&&tr[u].r<=y) return tr[u].sum;
  pushdown(u);
  return query(lc,x,y)+query(rc,x,y);
}
int query_path(int u,int v){ //查询路径
  int res=0;
  while(top[u]!=top[v]){
    if(dep[top[u]]<dep[top[v]]) swap(u,v);
    res+=query(1,id[top[u]],id[u]);
    u=fa[top[u]];
  }
  if(dep[u]<dep[v]) swap(u,v);
  res+=query(1,id[v],id[u]); //最后一段
  return res;
}
int query_tree(int u){ //查询子树
  return query(1,id[u],id[u]+siz[u]-1);
}
signed main(){
  ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
  cin>>n>>m>>root>>P;
  for(int i=1; i<=n; i++) cin>>w[i];
  for(int i=0; i<n-1; i++){
    cin>>a>>b;
    e[a].push_back(b); e[b].push_back(a);
  }
  dfs1(root,0); dfs2(root,root); //把树拆成链
  build(1,1,n);  //用链建线段树
  while(m--){
    cin>>t>>a;
    if(t==1) cin>>b>>c,change_path(a,b,c);
    else if(t==3) cin>>c,change_tree(a,c);
    else if(t==2) cin>>b,cout<<query_path(a,b)%P<<"\n";
    else cout<<query_tree(a)%P<<"\n";
  }
}