BZOJ3631 松鼠的新家(树上差分/点差分)

题目

n(n<=3e5)个点的一棵树，给出a1到an，

从a1走到a2，再从a2走到a3，从a3走到a4，...，从an-1走到an

要求，走的路径上每走一个点放一个糖果，最后an不放，

问每个点放了多少糖果

思路来源

https://www.luogu.org/problemnew/solution/P3258

题解

裸的树上点差分，a1走到a2，就把这一条路径+1即可

由于后续又把a2走到a3的路径+1，故需要对a2单点减一，减去重复

可知，a2到an这n-1个点都需减去重复

代码

#include<bits/stdc++.h> 
using namespace std;
typedef long long ll;
typedef pair<int,int> P;
#define pb push_back
const int N=3e5+10;
const int INF=0x3f3f3f3f;

int n,u,v,ans[N]; 
int a[N],dep[N],f[N][19];
vector<int>E[N];

void dfs(int u,int fa)
{
	for(int i=0;i<E[u].size();++i)
	{
		int v=E[u][i];
		if(v==fa)continue;
		dep[v]=dep[u]+1; //必须前序 否则传参 
		dfs(v,u);
		f[v][0]=u;
	}
}
void init()
{
	for(int i=1;i<19;++i)
	{
		for(int j=1;j<=n;++j)
		f[j][i]=f[f[j][i-1]][i-1];
	}
}
int lca(int x,int y)
{
	if(dep[x]<dep[y])swap(x,y);
	int d=dep[x]-dep[y];
	for(int i=18;i>=0;--i)
	if(d>>i&1)x=f[x][i];
	if(x==y)return x;
	for(int i=18;i>=0;--i)
	{
		if(f[x][i]!=f[y][i])
		x=f[x][i],y=f[y][i];
	}
	return f[x][0];
}
void add(int x,int y)//x到y的路径 +1
{
	ans[x]++;
	ans[y]++;
	int fa=lca(x,y);
	ans[fa]--;
	ans[f[fa][0]]--;
}
void dfs2(int u,int fa)
{
	for(int i=0;i<E[u].size();++i)
	{
		int v=E[u][i];
		if(v==fa)continue;
		dfs2(v,u);
		ans[u]+=ans[v];
	}
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<=n;++i)
	scanf("%d",&a[i]);
	for(int i=1;i<n;++i)
	{
		scanf("%d%d",&u,&v);
		E[v].pb(u);E[u].pb(v);
	}
	dfs(1,-1);init();
	for(int i=1;i<n;++i)
	add(a[i],a[i+1]);
	dfs2(1,-1);
	for(int i=2;i<=n;++i)
	ans[a[i]]--;
	for(int i=1;i<=n;++i)
	printf("%d\n",ans[i]);
    return 0;
}