洛谷3258 松鼠的新家

题目大意

按给定顺序遍历一棵n个节点的树，求每个节点经过的次数。（$2\le n\le 3\times 10^5$）

---

知识储备

树上差分：

首先回顾差分。差分就是将区间的修改转化为对两端重要节点的的修改。因此，树上差分的思想也就是将对路径的上所有节点的修改转化为对树上重要节点的修改。

设dif[u]为节点u在经过次数的差分数组中的值，我们知道，对差分数组求前缀和，结果就为原数组。我们的基本思路就为：对于每条路径，对路径上的重要节点进行修改，最后求出差分数组的前缀和，便是每个节点原本的经过次数。

树上的差分数组的前缀和其实就是以任意一节点为根的子树中所有节点的差分数组值之和。（联系下面的说明自己想想为什么）那么，任意一个节点被经过的次数，就等于以它为根的子树上所有节点的差分数组的值之和。

有了这个前提在先，下面对重要节点的差分数组值的修改就简单些了。

首先，从a到b的一条路径，必然会经过lca(a，b)，并经过a、b到lca(a，b)的路径上的所有祖先。因此a、b被经过了1次，那么它们各自路径上的所有祖先也被经过了1次。

但lca(a，b)既在a到lca(a，b)的路径上，也在b到lca(a，b)的路径上，会求和2次，而实际上lca(a，b)只被经过1次。因此dif[lca(a，b)]要减1。同时，lca(a，b)的父节点并没有被经过，但它求前缀和时会统计lca(a，b)的1次经过次数，因此lca(a，b)的父节点的经过次数也要减1。

综上所述，对于一条路径a->b，其差分数组的修改是这样的：
dif[a]++,dif[b]++,dif[lca(a，b)]--,dif[lca(a，b)的父节点]--

最后dfs递归求出其子树中所有差分数组的和，就是该节点的经过次数。

---

思路

树上差分模板题。求lca、修改、求和即可

---

代码

#include<iostream>
#include<cstdio>
#include<cmath>
#include<vector>
using namespace std;
const int N=1e6,P=20;
int n,x,y;
int dep[N],vis[N],f[N][P+1],dif[N];
vector <int> ed[N];
void dfs(int x,int fa)
{
	f[x][0]=fa;
	dep[x]=dep[fa]+1;
	for(int i=0;i<ed[x].size();i++)
	if(ed[x][i]!=fa) dfs(ed[x][i],x);
}
int query(int x)
{
	for(int i=0;i<ed[x].size();i++)
	if(ed[x][i]!=f[x][0])
	dif[x]+=query(ed[x][i]);
	return dif[x];
}
int lca(int x,int y)
{
	if(dep[x]>dep[y]) swap(x,y);
	int d=dep[y]-dep[x],step=0;
	while(d)
	{
		if(d&1) y=f[y][step];
		d>>=1,step++;
	}
	if(x==y) return x;
	step=P;
	while(--step>=0)
	if(f[x][step]!=f[y][step])
	x=f[x][step],y=f[y][step];
	return f[x][0];
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
	scanf("%d",&vis[i]);
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		ed[x].push_back(y),ed[y].push_back(x);
	}
	dfs(vis[1],0);
	for(int i=1;i<=P;i++)
	for(int j=1;j<=n;j++)
	f[j][i]=f[f[j][i-1]][i-1];
	for(int i=1;i<n;i++)//对每条路径进行修改 
	{
		int anc=lca(vis[i],vis[i+1]);
		dif[vis[i]]++,dif[vis[i+1]]++;
		dif[anc]--,dif[f[anc][0]]--;
	}
	query(1);//求差分数组的前缀和 
	dif[vis[1]]+=1;
	for(int i=1;i<=n;i++)
	printf("%d\n",dif[i]-1);//因为每个节点出去时也会统计一次（第一个除外，因此+1），经过次数都要-1 
	return 0;
}