轻重链剖分(树链剖分)模板

洛谷题目链接

/*
f数组表示当前节点的父节点
d数组表示当前节点的深度
size数组表示当前节点的子树大小
id当前节点的时间戳
rk当前节点时间戳对应当前节点的值(线段树的叶子节点是树上所有节点的时间戳构成的)
son当前节点的重儿子
top当前节点所在重链的顶端节点
*/
#include<bits/stdc++.h>
#define ls rt << 1
#define rs rt << 1|1
using namespace std;
const int maxn = 2e5 + 10;
struct node{
	int next,to;
}edge[maxn];

int tree[maxn << 2], add[maxn << 2];
int root, n, m;
int index0,a[maxn],mod,cnt, head[maxn], f[maxn],d[maxn],id[maxn],size[maxn],son[maxn],rk[maxn],top[maxn];
void add1(int from,int to){
	edge[cnt].to = to;
	edge[cnt].next = head[from];
	head[from] = cnt ++;
}
void pushdown(int l, int r ,int rt){
	if(add[rt]){
		int mid = (l + r) >> 1;
		add[rs] += add[rt];
		add[ls] += add[rt];
		tree[ls] += (mid - l + 1)*add[rt];
		tree[ls] %= mod;
		tree[rs] += (r - mid)*add[rt];
		tree[rs] %= mod;
		add[rt] = 0;
	}
}
void build(int l, int r, int rt){
	if(l == r){
		tree[rt] = rk[l]%mod;
		add[rt] = 0;
		return ;
	}

	int mid = (l + r) >> 1;
	build(l, mid , ls);
	build(mid + 1, r , rs);
	tree[rt] = tree[ls] + tree[rs];
}
void update(int L, int R ,int val, int l ,int r, int rt){
	if(L <= l && r <= R){
		tree[rt] = (tree[rt] + (r - l + 1)*val)%mod;
		add[rt] += val;
		return ;
	}

	int mid = (l + r) >> 1;
	pushdown(l, r, rt);
	if(mid >= L) update(L, R ,val ,l , mid, ls);
	if(mid < R) update(L, R , val, mid + 1 ,r ,rs);
	tree[rt] = tree[ls] + tree[rs];
}
int query(int L, int R, int l, int r ,int rt){
	if(L <= l && r <= R){
		return tree[rt];
	}
	pushdown(l , r, rt);
	int mid = (l + r) >> 1;
	int ans = 0;
	if(mid >= L) ans = (ans + query(L, R , l, mid , ls))%mod;
	if(mid < R) ans = (ans + query(L, R , mid + 1, r , rs))%mod;
	return ans;
}
void dfs1(int u, int fa, int depth){
	//cout << 1 << endl;
	f[u] = fa;
	d[u] = depth;
	size[u] = 1;
	for(int i = head[u]; ~i ; i = edge[i].next){
		int v = edge[i].to;
		if(v == fa) continue;
		dfs1(v, u, depth + 1);
		size[u] += size[v];
		if(size[v] > size[son[u]]){
			son[u] = v;
		}
	}
}
void dfs2(int u, int t){
	top[u] = t;
	id[u] = ++index0;
	rk[index0] = a[u];
	if(!son[u]) return;
	dfs2(son[u],t);

	for(int i = head[u]; ~i; i = edge[i].next){
		int v = edge[i].to;
		if(v != son[u] && v!= f[u]){
			dfs2(v, v);
		}
	}
}
void update1(int u, int v, int w){
	w %= mod;
	while(top[u] != top[v]){
		if(d[top[u]] < d[top[v]]){
			swap(u, v);
		}
		update(id[top[u]], id[u] , w, 1, n ,1);
		u = f[top[u]];
	}
	if(d[u] > d[v]){
		swap(u , v);
	}
	update(id[u], id[v], w, 1 , n , 1);
}

void query1(int u, int v){
	int ans = 0;
	while(top[u] != top[v]){
		if(d[top[u]] < d[top[v]]){
			swap(u, v);
		}
		ans = (ans + query(id[top[u]], id[u], 1, n, 1))%mod;
		u = f[top[u]];
	}
	if(d[u] > d[v]){
		swap(u, v);
	}
	ans = (ans + query(id[u], id[v], 1, n, 1))%mod;
	cout << ans << endl;
}
void updateson(int u, int k){
	k%= mod;
	update(id[u], id[u] + size[u] - 1, k, 1, n , 1);
}
void queryson(int u){
	//cout << id[u] <<" " << id[u] + size[u] - 1 << endl;
	cout << query(id[u], id[u] + size[u] - 1, 1, n , 1) %mod << endl;
}
int main()
{
	freopen("in.txt","r",stdin);
	freopen("out.txt","w",stdout);
	memset(head, -1, sizeof head);
	cin  >> n >> m>> root>> mod;
	for(int i = 1; i <= n; i ++){
		scanf("%d",&a[i]);
	}
	for(int i = 0; i < n -  1;i ++){
		int u, v;
		cin >> u >> v;
		add1(u, v);
		add1(v, u);
	}

	dfs1(root, 0 , 1);
	
	dfs2(root, root);

	build(1, n , 1);
	for(int i = 0 ; i < m; i ++){
		int x;
		cin >> x;
		if(x == 1){
			int u,v,w;
			cin >> u >> v >> w;
			update1(u, v, w);
		}
		if(x == 2){
			int u,v;
			cin >> u >> v;
			query1(u, v);
		}
		if(x == 3){
			int u, v;
			cin >> u>> v;
			updateson(u, v);
		}
		if(x == 4){
			int u;
			cin >> u;
			queryson(u);
		}

	}
}