摘自博客:https://blog.csdn.net/VictoryCzt/article/details/85268376
博主讲得十分的清楚
以下是带有注释的代码:
1 /* 2 这是一份自用的树链剖分模板, 3 自己加的注释,便于更好的理解 4 该模板支持的操作: 5 1.换根 6 2.一条链上点权加 7 3.一个子树内点权加 8 4.询问一条链上点权的和 9 5.询问一个子树内的点权和 10 6.求LCA 11 */ 12 #include<bits/stdc++.h> 13 #define numm ch-48 14 #define pd putchar(' ') 15 #define pn putchar(' ') 16 #define pb push_back 17 #define debug(args...) cout<<#args<<"->"<<args<<endl 18 #define bug cout<<"************" 19 using namespace std; 20 template <typename T> 21 void read(T &res) { 22 bool flag=false;char ch; 23 while(!isdigit(ch=getchar())) (ch=='-')&&(flag=true); 24 for(res=numm;isdigit(ch=getchar());res=(res<<1)+(res<<3)+numm); 25 flag&&(res=-res); 26 } 27 template <typename T> 28 void write(T x) { 29 if(x<0) putchar('-'),x=-x; 30 if(x>9) write(x/10); 31 putchar(x%10+'0'); 32 } 33 typedef long long ll; 34 const int inf=0x3f3f3f3f; 35 const int maxn=1e5+10; 36 /*建图部分*/ 37 struct node {///链式前向星 38 int v,net; 39 node(){} 40 node(int v,int net):v(v),net(net){} 41 }e[maxn<<1];///树有双向边,边数*2 42 int head[maxn],cnt,n; 43 void init() { 44 cnt=0; 45 for(int i=1;i<=n;i++) 46 head[i]=-1; 47 } 48 void addedge(int u,int v) {///建双向边 49 e[cnt]=node(v,head[u]); 50 head[u]=cnt++; 51 e[cnt]=node(u,head[v]); 52 head[v]=cnt++; 53 } 54 int id[maxn],rk[maxn],fa[maxn],top[maxn],dep[maxn],son[maxn],sze[maxn],val[maxn],root=1; 55 //id[u]:点u的dfs序,一条重链上的点标号连续 56 //sze[u]:点u的子树结点个数,包含点u 57 //son[u]:点u的重儿子的id 58 //fa[u]:点u的父亲,根结点1的父亲为0 59 //top[u]:点u所在重链的头结点,轻链则为本身,则top[u]=u; 60 //val[u]:点u的初始权值 61 //rk[id[u]]:id[u]对应的u结点 62 //dep[u]:dfs时,u对应的深度 63 int minp[maxn<<2]; 64 ll sum[maxn<<2],lazy[maxn<<2]; 65 /*线段树部分*/ 66 //sum[k]:点k所在区间的树结点总和 67 //lazy[k]:点k所在区间的树结点总和的懒标记 68 //minp[k]:点k所在区间内的最小深度的点u的下标 69 int get_minp(int a,int b) { ///返回a,b结点中深度最小的结点 70 if(a==inf) return b; 71 if(b==inf) return a; 72 if(dep[a]<dep[b]) return a; 73 return b; 74 } 75 void pushup(int k) { 76 sum[k]=sum[k<<1]+sum[k<<1|1]; 77 minp[k]=get_minp(minp[k<<1],minp[k<<1|1]); 78 } 79 void pushdown(int k,int l,int r,int mid) { 80 if(!lazy[k]) return ; 81 sum[k<<1]+=(mid-l+1)*lazy[k]; 82 sum[k<<1|1]+=(r-mid)*lazy[k]; 83 lazy[k<<1]+=lazy[k]; 84 lazy[k<<1|1]+=lazy[k]; 85 lazy[k]=0; 86 } 87 void build(int k,int l,int r) {///建树 88 if(l==r) { 89 sum[k]=val[rk[l]]; 90 minp[k]=rk[l]; 91 return ; 92 } 93 int mid=l+r>>1; 94 build(k<<1,l,mid); 95 build(k<<1|1,mid+1,r); 96 pushup(k); 97 } 98 void update(int k,int l,int r,int ql,int qr,ll w) {///区间修改 99 if(ql<=l&&r<=qr) { 100 sum[k]+=(r-l+1)*w; 101 lazy[k]+=w; 102 return ; 103 } 104 int mid=l+r>>1; 105 pushdown(k,l,r,mid);///attention!!! 106 if(ql<=mid) update(k<<1,l,mid,ql,qr,w); 107 if(mid<qr) update(k<<1|1,mid+1,r,ql,qr,w); 108 pushup(k); 109 } 110 int query_pos(int k,int l,int r,int ql,int qr) {///区间查询最小深度的结点 111 if(ql>qr) return inf; 112 if(ql<=l&&r<=qr) return minp[k]; 113 int mid=l+r>>1; 114 if(qr<=mid) return query_pos(k<<1,l,mid,ql,qr); 115 else if(ql>mid) return query_pos(k<<1|1,mid+1,r,ql,qr); 116 return get_minp(query_pos(k<<1,l,mid,ql,qr),query_pos(k<<1|1,mid+1,r,ql,qr)); 117 } 118 ll query_sum(int k,int l,int r,int ql,int qr) {///区间求和 119 if(ql<=l&&r<=qr) return sum[k]; 120 int mid=l+r>>1; 121 pushdown(k,l,r,mid); 122 if(qr<=mid) return query_sum(k<<1,l,mid,ql,qr); 123 else if(ql>mid) return query_sum(k<<1|1,mid+1,r,ql,qr); 124 else return query_sum(k<<1,l,mid,ql,qr)+query_sum(k<<1|1,mid+1,r,ql,qr); 125 } 126 /*树剖部分*/ 127 void dfs1(int u) { ///得到每个点的深度,父亲,,子树大小和重儿子 128 sze[u]=1; ///把自己算进去 129 for(int i=head[u];~i;i=e[i].net) { 130 int v=e[i].v; 131 if(v==fa[u]) continue; 132 fa[v]=u; 133 dep[v]=dep[u]+1; 134 dfs1(v); 135 sze[u]+=sze[v]; ///回溯累加 136 if(!son[u]||sze[son[u]]<sze[v])///更新重儿子 137 son[u]=v; 138 } 139 } 140 void dfs2(int u,int root) { 141 top[u]=root; 142 id[u]=++cnt; 143 rk[cnt]=u; 144 if(!son[u]) return ; 145 dfs2(son[u],root); 146 ///先dfs重儿子,使得每条重链上的结点标号连续 147 for(int i=head[u];~i;i=e[i].net) { 148 int v=e[i].v; 149 if(v==son[u]||v==fa[u]) continue; 150 dfs2(v,v);///轻链的top则为本身 151 } 152 } 153 int lca(int a,int b) { ///求公共祖先 154 while(top[a]!=top[b]) { 155 if(dep[top[a]]<dep[top[b]]) swap(a,b); 156 a=fa[top[a]]; 157 } 158 if(dep[a]>dep[b]) swap(a,b); 159 return a; 160 } 161 int Upto(int a,int b) { ///返回a到b上深度最小的点的下标 162 int ans=inf; 163 while(top[a]!=top[b]) { 164 if(dep[top[a]]<dep[top[b]]) swap(a,b); 165 ///不能包含u点,要分类讨论 166 if(top[a]!=b) ans=get_minp(ans,query_pos(1,1,n,id[top[a]],id[a])); 167 else ans=get_minp(ans,query_pos(1,1,n,id[top[a]]+1,id[a])); 168 a=fa[top[a]]; 169 } 170 if(dep[a]>dep[b]) swap(a,b); 171 172 ans=get_minp(ans,query_pos(1,1,n,id[a]+1,id[b]));///不能包含u点,所以+1 173 return ans; 174 } 175 int check(int u) { ///判r与root的关系 176 if(u==root) return 0;///此时操作范围就是整个树,将整个树加或者求和即可。 177 int L=lca(u,root); 178 if(L==u) return Upto(u,root);/*当r在u的子树内时,操作范围就为 179 整个树减去u到r路径上深度最小的点(不包含u点)的原来的子树*/ 180 else return -1;///当r不在u的原来的子树里面时,操作范围就是原来的子树。 181 } 182 void add_chain(int a,int b,ll w) { ///一条链上点权加 183 while(top[a]!=top[b]) { 184 if(dep[top[a]]<dep[top[b]]) swap(a,b);///找到深度最大的点 185 update(1,1,n,id[top[a]],id[a],w); 186 a=fa[top[a]]; ///一次跳一个 187 } 188 if(dep[a]>dep[b]) swap(a,b); 189 update(1,1,n,id[a],id[b],w); 190 } 191 ll ask_chain(int a,int b) { ///询问一条链上点权的和 192 ll ans=0; 193 while(top[a]!=top[b]) { 194 if(dep[top[a]]<dep[top[b]]) swap(a,b); 195 ans+=query_sum(1,1,n,id[top[a]],id[a]); 196 a=fa[top[a]]; 197 } 198 if(dep[a]>dep[b]) swap(a,b); 199 ans+=query_sum(1,1,n,id[a],id[b]); 200 return ans; 201 } 202 void add_tree(int u,ll w) { ///一个子树内点权加 203 int type=check(u); ///判断u与root的情况 204 if(!type) update(1,1,n,1,n,w); 205 else if(type>0) { 206 update(1,1,n,1,n,w); 207 update(1,1,n,id[type],id[type]+sze[type]-1,-w); 208 } 209 else update(1,1,n,id[u],id[u]+sze[u]-1,w); 210 } 211 ll ask_tree(int u) { ///询问u的子树内的点权和 212 int type=check(u); ///判断u与root的情况 213 if(!type) return query_sum(1,1,n,1,n); 214 else if(type>0){ 215 ll ans=query_sum(1,1,n,id[type],id[type]+sze[type]-1); 216 return query_sum(1,1,n,1,n)-ans; 217 } 218 else return query_sum(1,1,n,id[u],id[u]+sze[u]-1); 219 } 220 int main() 221 { 222 int a,b; 223 ll w; 224 read(n); 225 init(); 226 for(int i=1;i<=n;i++) ///得到每个点的初始权值 227 read(val[i]); 228 for(int i=2;i<=n;i++) { ///得到每个点的父亲 229 read(a); 230 addedge(a,i); ///建边 231 } 232 dfs1(root); 233 cnt=0; ///attention!!! 234 dfs2(1,1); 235 build(1,1,n); 236 int m; 237 read(m); 238 while(m--) { 239 int oper; 240 read(oper); 241 if(oper==1) { ///换根 242 read(a); 243 root=a; 244 } 245 else if(oper==2) { ///一条链上点权加 246 read(a),read(b),read(w); 247 add_chain(a,b,w); 248 } 249 else if(oper==3) { ///一个子树内点权加 250 read(a),read(w); 251 add_tree(a,w); 252 } 253 else if(oper==4) { ///询问一条链上点权的和 254 read(a),read(b); 255 write(ask_chain(a,b));pn; 256 } 257 else if(oper==5) { ///询问一个子树内的点权和 258 read(a); 259 write(ask_tree(a));pn; 260 } 261 } 262 return 0; 263 }