树链剖分模版

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

猜你喜欢

转载自www.cnblogs.com/-Ackerman/p/11391343.html