树链剖分模版

  1 #include <stdio.h>
  2 #include <cstring>
  3 #include <iostream>
  4 #include <string>
  5 #include <algorithm>
  6 #include <queue>
  7 #include <vector>
  8 #include <math.h>
  9 #include <map>
 10 
 11 #define LL long long
 12 using namespace std;
 13 const int maxn = 2e5 + 10;
 14 
 15 struct Edge{
 16     int to,next;
 17 }edge[maxn];
 18 
 19 int tot,head[maxn];
 20 
 21 void add_edge(int u,int v){
 22     edge[++tot] = Edge{v,head[u]};
 23     head[u] = tot;
 24 }
 25 
 26 int mod;
 27 int v[maxn]; // 结点权值
 28 int fa[maxn]; // 父亲
 29 int dep[maxn]; // 深度
 30 int siz[maxn]; // 大小
 31 int son[maxn]; // 重儿子
 32 
 33 // 第一遍的dfs可以得到深度，父亲，大小，重儿子
 34 void dfs1(int u,int f){ // u是当前结点，f是u的父亲
 35     fa[u] = f;
 36     dep[u] = dep[f] + 1;
 37     siz[u] = 1;
 38     int maxsize = -1;  // 判断是不是重儿子的一个临时变量
 39     for (int i=head[u];~i;i=edge[i].next){
 40         int v = edge[i].to;
 41         if (v == f)   //如果是父亲肯定不可以
 42             continue;
 43         dfs1(v,u);
 44         siz[u] += siz[v];
 45         if (siz[v] > maxsize){
 46             maxsize = siz[v];
 47             son[u] = v;  // u的重儿子就是v
 48         }
 49     }
 50 }
 51 
 52 int tim; // 时间戳计数器
 53 int dfn[maxn]; // 时间戳
 54 int top[maxn]; //重链的顶部
 55 int w[maxn]; // 结点权值dfs序
 56 
 57 void dfs2(int u,int t){ // u是当前结点，t是当前结点的重链的头
 58     dfn[u] = ++tim;
 59     top[u] = t;
 60     w[tim] = v[u];
 61     if (!son[u])  // 如果不存在重儿子，那么它就是叶子结点，退出
 62         return;
 63     dfs2(son[u],t);
 64     for (int i=head[u];~i;i=edge[i].next){
 65         int v = edge[i].to;
 66         if (v == fa[u] || v == son[u])   // 往上遍历肯定不可以了，因为前面的dfs2已经遍历了重儿子所以这里也没必要了
 67             continue;
 68         dfs2(v,v);   // 此时这个肯定是轻儿子
 69     }
 70 }
 71 
 72 struct segment_tree{
 73     int l,r,val;
 74     int lazy;
 75 }tree[maxn*4];
 76 
 77 void pushup(int nod){
 78     tree[nod].val = (tree[nod<<1].val + tree[(nod<<1)+1].val) % mod;
 79 }
 80 
 81 void pushdown(int nod){
 82     tree[nod<<1].lazy += tree[nod].lazy;
 83     tree[(nod<<1)+1].lazy += tree[nod].lazy;
 84     tree[nod<<1].val += (tree[nod<<1].r-tree[nod<<1].l + 1) * tree[nod].lazy % mod;
 85     tree[(nod<<1)+1].val += (tree[(nod<<1)+1].r-tree[(nod<<1)+1].l+1) * tree[nod].lazy % mod;
 86     tree[nod].lazy = 0;
 87 }
 88 
 89 void build(int l,int r,int nod=1){
 90     tree[nod].l = l;
 91     tree[nod].r = r;
 92     if (l == r){
 93         tree[nod].lazy = 0;
 94         tree[nod].val = w[l] % mod;
 95         return ;
 96     }
 97     int mid = (l+r)>>1;
 98     build(l,mid,nod<<1);
 99     build(mid+1,r,(nod<<1)+1);
100     pushup(nod);
101 }
102 
103 void modify(int x,int y,int z,int k=1){
104     int l = tree[k].l, r = tree[k].r;
105     if (x<= l && y>=r){
106         tree[k].lazy += z;
107         tree[k].val += (r-l+1) * z;
108         tree[k].val %= mod;
109         return ;
110     }
111     if (tree[k].lazy)
112         pushdown(k);
113     int mid = (l+r)>>1;
114     if (x<=mid){
115         modify(x,y,z,k<<1);
116     }
117     if (y>mid){
118         modify(x,y,z,(k<<1)+1);
119     }
120     pushup(k);
121 }
122 
123 int query(int x,int y,int k=1){
124     int l = tree[k].l,r = tree[k].r;
125     if (x<=l && y>=r){
126         return tree[k].val;
127     }
128     if (tree[k].lazy){
129         pushdown(k);
130     }
131     int sum = 0,mid = (l+r)>>1;
132     if (x <= mid){
133         sum += query(x,y,k<<1);
134     }
135     if (y > mid){
136         sum += query(x,y,(k<<1)+1);
137     }
138     return sum % mod;
139 }
140 
141 void mchain(int x,int y,int z){  // 结点x->结点y 最短路径上所有结点加z
142     z %= mod;
143     while (top[x] != top[y]){
144         if (dep[top[x]] < dep[top[y]])
145             swap(x,y);
146         modify(dfn[top[x]],dfn[x],z);
147         x = fa[top[x]];
148     }
149     if (dep[x] > dep[y])
150         swap(x,y);
151     modify(dfn[x],dfn[y],z);
152 }
153 
154 int qchain(int x,int y){  // 查询x到y结点最短路径上所有节点的值之和
155     int ret = 0;
156     while (top[x] != top[y]){
157         if (dep[top[x]] < dep[top[y]])
158             swap(x,y);
159         ret += query(dfn[top[x]],dfn[x]);
160         x = fa[top[x]];
161     }
162     if (dep[x] > dep[y])
163         swap(x,y);
164     ret += query(dfn[x],dfn[y]);
165     return ret % mod;
166 }
167 
168 void mson(int x,int z){ // 以x为根节点的子树内所有节点值都加上z
169     modify(dfn[x],dfn[x]+siz[x]-1,z);  // 必定是连续的
170 }
171 
172 int qson(int x){  // 以x为根节点的子树内所有节点值之和
173     return query(dfn[x],dfn[x]+siz[x]-1);
174 }
175 
176 
177 int main(){
178     memset(head,-1, sizeof(head));
179     int n,m,r;
180     scanf("%d%d%d%d",&n,&m,&r,&mod);
181     for (int i=1;i<=n;i++){
182         scanf("%d",&v[i]);
183     }
184     for (int i=1;i<n;i++){
185         int u,v;
186         scanf("%d%d",&u,&v);
187         add_edge(u,v);
188         add_edge(v,u);
189     }
190     dfs1(r,r);
191     dfs2(r,r);
192     build(1,n);
193     while (m--){
194         int opt,x,y,z;
195         scanf("%d",&opt);
196         switch(opt){
197             case 1:
198                 scanf("%d%d%d",&x,&y,&z);
199                 mchain(x,y,z);
200                 break;
201             case 2:
202                 scanf("%d%d",&x,&y);
203                 printf("%d\n",qchain(x,y));
204                 break;
205             case 3:
206                 scanf("%d%d",&x,&z);
207                 mson(x,z);
208                 break;
209             case 4:
210                 scanf("%d",&x);
211                 printf("%d\n",qson(x));
212                 break;
213         }
214     }
215     return 0;
216 }