洛谷 [P3258] 松鼠的新家

树上差分

对于一条路径 \(u->v\) 来说,设 \(t=LCA(u,v)\) ,d[]为差分数组 ,则有 d[u]++;d[v]++;d[t]--;d[fa[t]]--;

注意:题目中所给的路径上的点都多计算了一次,统计答案时要减去

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int MAXN=3000005;
int num[MAXN],fa[MAXN],ans[MAXN],d[MAXN],head[MAXN],dep[MAXN],siz[MAXN],son[MAXN],top[MAXN],n,m,nume;
struct edge{
	int to,nxt;
}e[MAXN<<1];
void adde(int from,int to){
	e[++nume].to=to;
	e[nume].nxt=head[from];
	head[from]=nume;
}
int init(){
	int rv=0,fh=1;
	char c=getchar();
	while(c<'0'||c>'9'){
		if(c=='-') fh=-1;
		c=getchar();
	}
	while(c>='0'&&c<='9'){
		rv=(rv<<1)+(rv<<3)+c-'0';
		c=getchar();
	}
	return fh*rv;
}
void dfs1(int u,int rt){
	fa[u]=rt;
	siz[u]=1;
	dep[u]=dep[rt]+1;
	int ma=0;
	for(int i=head[u];i;i=e[i].nxt){
		int v=e[i].to;
		if(v==rt) continue;
		dfs1(v,u);
		siz[u]+=siz[v];
		if(ma<siz[v]){
			ma=siz[u];
			son[u]=v;
		}
	}
}
void dfs2(int u,int topf){
	top[u]=topf;
	if(!son[u]) return;
	dfs2(son[u],topf);
	for(int i=head[u];i;i=e[i].nxt){
		int v=e[i].to;
		if(v==fa[u]||v==son[u]) continue;
		dfs2(v,v);
	}
}
int LCA(int a,int b){
	while(top[a]!=top[b]){
		if(dep[top[a]]>dep[top[b]]) a=fa[top[a]];
		else if(dep[top[a]]<dep[top[b]]) b=fa[top[b]];
		else b=fa[top[b]],a=fa[top[a]];
	}
	return dep[b]<dep[a]?b:a;
}
void work(int u){
	ans[u]=d[u];
	for(int i=head[u];i;i=e[i].nxt){
		int v=e[i].to;
		if(v==fa[u]) continue;
		work(v);
		ans[u]+=ans[v];
	}
}
int main(){
	n=init();
	for(int i=1;i<=n;i++) num[i]=init();
	for(int i=1;i<n;i++){
		int u=init(),v=init();
		adde(u,v);adde(v,u);
	}
	dep[1]=1;
	dfs1(1,0);
	dfs2(1,1);
	for(int i=1;i<n;i++){
		int t=LCA(num[i],num[i+1]);
		//cout<<t<<endl;
		d[num[i]]++;
		d[num[i+1]]++;
		d[t]--;d[fa[t]]--;
	}
	work(1);
	for(int i=2;i<=n;i++) ans[num[i]]--;
	for(int i=1;i<=n;i++) cout<<ans[i]<<endl;
	return 0;
}