洛谷P3258 [JLOI]2014松鼠的新家

题目

树上差分

树上点差分，注意会出现路径端点多记录的情况，这时需要在最后输出的时候输出子树的差分数组的和-1，而不是在处理原数据的时候减1。并且a[n]不需要糖果，最后也减去就行。

#include <bits/stdc++.h>
#define N 1001011
using namespace std;
struct edg {
	int to, nex;
}e[N];
int n, a[N], cnt, lin[N], fa[N][19], dep[N], sum[N];//sum[i]乃差分数组。
inline void add(int f, int t)
{
	e[++cnt].to = t;
	e[cnt].nex = lin[f];
	lin[f] = cnt;
}
void dfs(int now, int f)
{
 	fa[now][0] = f;
 	dep[now] = dep[f] + 1;
 	for (int i = lin[now]; i; i = e[i].nex)
 	{
 		int to = e[i].to;
	 	if (to == f) continue;
 		dfs(to, now);
 	}
}
int dfs2(int now)
{
	for (int i = lin[now]; i; i = e[i].nex)
	{
		int to = e[i].to;
		if (to == fa[now][0]) continue;
		dfs2(to);
		sum[now] += sum[to];
	}
	return sum[now];
}
inline void init()
{
 	dfs(1, 0);
 	for (int j = 1; j <= 18; j++)
 		for (int i = 1; i <= n; i++)
 			fa[i][j] = fa[fa[i][j - 1]][j - 1];
}
int lca(int u, int v)
{
	if (dep[u] > dep[v])
		swap(u, v);
	for (int k = 0; k <= 18; k++)
		if ((dep[v] - dep[u]) >> k & 1)
			v = fa[v][k];
	if (u == v) return u;
	for (int k = 18; k >= 0; k--)
		if (fa[u][k] != fa[v][k])
			u = fa[u][k], v = fa[v][k];
	return fa[u][0];
}
int main()
{
	scanf("%d", &n);
	for (int i = 1; i <= n; i++)
		scanf("%d", &a[i]);
	for (int i = 1, u, v; i < n; i++)
	{
		scanf("%d%d", &u, &v);
		add(u, v);
		add(v, u);
	}
	init();
	for (int i = 2; i <= n; i++)
	{
		int now = lca(a[i], a[i - 1]);//now是当前经过的两点的lca。
		sum[a[i]]++, sum[a[i - 1]]++, sum[now]--, sum[fa[now][0]]--;
	}
	dfs2(1);
	for (int i = 2; i <= n; i++)
		sum[a[i]]--;
	for (int i = 1; i <= n; i++)
		printf("%d\n", sum[i]);
	return 0;
}