洛谷P3178 [HAOI2015]树上操作（树链剖分）

P3178 [HAOI2015]树上操作

题目描述
有一棵点数为 N 的树，以点 1 为根，且树点有边权。然后有 M 个操作，分为三种：操作 1 ：把某个节点 x 的点权增加 a 。操作 2 ：把某个节点 x 为根的子树中所有点的点权都增加 a 。操作 3 ：询问某个节点 x 到根的路径中所有点的点权和。

输入输出格式
输入格式：
第一行包含两个整数 N, M 。表示点数和操作数。接下来一行 N 个整数，表示树中节点的初始权值。接下来 N-1 行每行两个正整数 from, to ， 表示该树中存在一条边 (from, to) 。再接下来 M 行，每行分别表示一次操作。其中第一个数表示该操作的种类（ 1-3 ） ，之后接这个操作的参数（ x 或者 x a ） 。

输出格式：
对于每个询问操作，输出该询问的答案。答案之间用换行隔开。

输入输出样例

输入样例#1：
5 5
1 2 3 4 5
1 2
1 4
2 3
2 5
3 3
1 2 1
3 5
2 1 2
3 3
输出样例#1：
6
9
13
说明
对于 100% 的数据， N,M<=100000 ，且所有输入数据的绝对值都不会超过 10^6 。

代码

#include <cstdio>
#include <algorithm>
#define ll long long
#define N 300000
using namespace std;

struct arr
{
	int to,nxt;
}e[N];
struct tree
{
	int l,r;
	ll sm,lz;
}tr[N];
int n,m;
ll a[N],b[N];
int ls[N],l,cnt;
int dep[N],son[N],fa[N],siz[N];
int id[N],top[N];

void update(int p){tr[p].sm = tr[p * 2].sm + tr[p * 2 + 1].sm;}

void add(int x, int y)
{
	e[++l].to = y;
	e[l].nxt = ls[x];
	ls[x] = l;
}

void dfs1(int x, int father, int deep)
{
	siz[x] = 1;
	fa[x] = father;
	dep[x] = deep;
	for (int i = ls[x]; i; i = e[i].nxt)
		if (e[i].to != father)
		{
			dfs1(e[i].to, x, deep + 1);
			siz[x] += siz[e[i].to];
			if (!son[x] || siz[e[i].to] > siz[son[x]]) son[x] = e[i].to;
		}
}

void dfs2(int x, int tp)
{
	id[x] = ++cnt;
	top[x] = tp;
	a[cnt] = b[x];
	if (!son[x]) return;
	dfs2(son[x], tp);
	for (int i = ls[x]; i; i = e[i].nxt)
		if (!id[e[i].to]) dfs2(e[i].to, e[i].to);
}

void build(int p, int l, int r)
{
	tr[p].l = l;
	tr[p].r = r;
	if (l == r)
	{
		tr[p].sm = a[l];
		return;
	}
	int mid = (tr[p].l + tr[p].r) / 2;
	build(p * 2, l, mid);
	build(p * 2 + 1, mid + 1, r);
	update(p);
}

void down(int p)
{
	if (!tr[p].lz) return;
	int ls = p * 2;
	int rs = ls + 1;
	tr[ls].sm += tr[p].lz * (tr[ls].r - tr[ls].l + 1);
	tr[rs].sm += tr[p].lz * (tr[rs].r - tr[rs].l + 1);
	tr[ls].lz += tr[p].lz;
	tr[rs].lz += tr[p].lz;
	tr[p].lz = 0;
}

void ins(int p, int x, int y, ll w)
{
	if (x <= tr[p].l && y >= tr[p].r)
	{
		tr[p].sm += (tr[p].r - tr[p].l + 1) * w;
		tr[p].lz += w;
		return;
	}
	down(p);
	int mid = (tr[p].l + tr[p].r) / 2;
	if (x <= mid) ins(p * 2, x, y, w);
	if (y > mid) ins(p * 2 + 1, x, y, w);
	update(p);
}

ll sum(int p, int x, int y)
{
	ll ans = 0;
	if (x <= tr[p].l && y >= tr[p].r) return tr[p].sm;
	down(p);
	int mid = (tr[p].l + tr[p].r) / 2;
	if (x <= mid) ans += sum(p * 2, x, y);
	if (y > mid) ans += sum(p * 2 + 1, x, y);
	return ans;
}

ll treesum(int x, int y)
{
	ll ans = 0;
	while (top[x] != top[y])
	{
		if (dep[top[x]] < dep[top[y]]) swap(x, y);
		ans += sum(1, id[top[x]], id[x]);
		x = fa[top[x]];
	}
	if (dep[x] > dep[y]) swap(x, y);
	ans += sum(1, id[x], id[y]);
	return ans;
}

int main()
{
	scanf("%d%d", &n, &m);
	for (int i = 1; i <= n; i++) scanf("%lld", &b[i]);
	for (int i = 1; i < n; i++)
	{
		int x, y;
		scanf("%d%d", &x ,&y);
		add(x, y);
		add(y, x);
	}
	dfs1(1, 0, 1);
	dfs2(1, 1);
	build(1, 1, n);
	while (m--)
	{
		int o, x;
		ll w;
		scanf("%d", &o);
		if (o == 1)
		{
			scanf("%d%lld", &x, &w);
			ins(1, id[x], id[x], w);
		}
		if (o == 2)
		{
			scanf("%d%lld", &x, &w);
			ins(1, id[x], id[x] + siz[x] - 1, w);
		}
		if (o == 3)
		{
			scanf("%d", &x);
			printf("%lld\n", treesum(1, x));
		}
	}
}