树链剖分

模板题

题目

题目描述

如题，已知一棵包含 $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$ 取模）。

代码

#include <iostream>
#include <cstdio>

#define lson l, mid, rt << 1
#define rson mid + 1, r, rt << 1 | 1

const int MAXN = 100005;
int n, m, r, p;
int tot, head[MAXN];
int a[MAXN];
int cnt, id[MAXN], w[MAXN], sum[MAXN << 2], col[MAXN << 2];
int deep[MAXN], fa[MAXN], top[MAXN];
int size[MAXN], son[MAXN];
struct Edge
{
	int from, to, next;

	Edge(int from = 0, int to = 0, int next = 0, int cost = 0) : from(from), to(to), next(next) {}
} edge[MAXN << 1];

void Addedge(int u, int v)
{
	edge[++ tot] = Edge(u, v, head[u]);

	head[u] = tot;
}

// 树剖预处理
void Dfs1(int u, int f)
{
	deep[u] = deep[f] + 1;
	fa[u] = f;

	int maxson = -1;
	size[u] = 1;

	for(int i = head[u]; i; i = edge[i].next)
	{
		int v = edge[i].to;

		if(v == f) continue;

		Dfs1(v, u);

		if(size[v] > maxson)
			son[u] = v, maxson = size[v];

		size[u] += size[v];
	}
}

void Dfs2(int u, int t)
{
	id[u] = ++ cnt;
	w[cnt] = a[u];
	top[u] = t;

	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 == fa[u]) continue;

		Dfs2(v, v);
	}
}
//

// 线段树
void Update(int rt)
{
	sum[rt] = (sum[rt << 1] + sum[rt << 1 | 1]) % p;
}

void Build(int l, int r, int rt)
{
	if(l == r)
	{
		sum[rt] = w[l];

		return;
	}

	int mid = (l + r) >> 1;

	Build(lson);
	Build(rson);

	Update(rt);
}

void Color(int l, int r, int rt, int v)
{
	sum[rt] = (sum[rt] + (r - l + 1) * v) % p;

	col[rt] = (col[rt] + v) % p;
}

void Push_col(int l, int r, int rt)
{
	if(col[rt])
	{
		int mid = (l + r) >> 1;

		Color(lson, col[rt]);
		Color(rson, col[rt]);

		col[rt] = 0;
	}
}

void Modify(int l, int r, int rt, int ql, int qr, int v)
{
	if(l >= ql && r <= qr)
	{
		Color(l, r, rt, v);

		return;
	}

	Push_col(l, r, rt);

	int mid = (l + r) >> 1;

	if(ql <= mid) Modify(lson, ql, qr, v);
	if(qr > mid) Modify(rson, ql, qr, v);

	Update(rt);
}

int Quiry(int l, int r, int rt, int ql, int qr)
{
	if(l >= ql && r <= qr) {return sum[rt];}

	Push_col(l, r, rt);

	int mid = (l + r) >> 1;

	if(ql <= mid)
	{
		if(qr <= mid)
			return Quiry(lson, ql, qr);
		else
			return (Quiry(lson, ql, qr) + Quiry(rson, ql, qr)) % p;
	}
	else
		return Quiry(rson, ql, qr);
}
//

// 树剖
void Mrange(int x, int y, int v)
{
	while(top[x] != top[y])
	{
		if(deep[top[x]] > deep[top[y]]) std:: swap(x, y);

		Modify(1, n, 1, id[top[y]], id[y], v);

		y = fa[top[y]];
	}

	if(deep[x] > deep[y]) std:: swap(x, y);

	Modify(1, n, 1, id[x], id[y], v);
}

void Mson(int x, int v)
{
	Modify(1, n, 1, id[x], id[x] + size[x] - 1, v);
}

int Qrange(int x, int y)
{
	int res = 0;

	while(top[x] != top[y])
	{
		if(deep[top[x]] > deep[top[y]]) std:: swap(x, y);

		res = (res + Quiry(1, n, 1, id[top[y]], id[y])) % p;

		y = fa[top[y]];
	}

	if(deep[x] > deep[y]) std:: swap(x, y);

	res = (res + Quiry(1, n, 1, id[x], id[y])) % p;

	return res;
}

int Qson(int x)
{
	return Quiry(1, n, 1, id[x], id[x] + size[x] - 1);
}
//

int main()
{
	scanf("%d%d%d%d", &n, &m, &r, &p);
	for(int i = 1; i <= n; i ++) scanf("%d", &a[i]);
	for(int i = 1; i < n; i ++)
	{
		int x, y;

		scanf("%d%d", &x, &y);

		Addedge(x, y);
		Addedge(y, x);
	}

	Dfs1(r, 0);
	Dfs2(r, r);
	Build(1, n, 1);

	for(int i = 1; i <= m; i ++)
	{	
		int c, x, y, z;
		
		scanf("%d", &c);
		
		if(c == 1)
		{
			scanf("%d%d%d", &x, &y, &z);
			
			Mrange(x, y, z);
			
			continue;
		}
		if(c == 2)
		{
			scanf("%d%d", &x, &y);
			
			printf("%d\n", Qrange(x, y));
			
			continue;
		}
		if(c == 3)
		{
			scanf("%d%d", &x, &z);
			
			Mson(x, z);
			
			continue;
		}
		if(c == 4)
		{
			scanf("%d", &x);
			
			printf("%d\n", Qson(x));
		}
	}
	
	return 0;
}