[洛谷2664] 樹上游戲

傳送門

看起來像是點分，其實有 $O(n)$ 做法？
其實我就是不想寫點分才這麼寫結果感覺腦子燒掉了

考慮一種顏色對答案的貢獻。

考慮把樹中這種顏色的點都刪掉，那麼就會有很多的小樹，這些小樹中的點互相之間不會產生貢獻，而不同樹的兩個點之間會產生貢獻。

由此，我們可以得到每一種顏色，點的sum值就是 n - 所在小樹的size。

由此，一個點總的sum就是 n * 顏色數 - 每種顏色時所在小樹的size。

我們考慮對於一棵小樹的size，存在深度最小的節點上，那麼後面就可以用樹上差分實現覆蓋。

嗯，大概用虛樹是可以的，可是我不會
所以size怎麼求？我也不會講，亂搞就好了

#include<bits/stdc++.h>
#define LL long long
#define re register
#define fr(i,x,y) for(int i=(x);i<=(y);i++)
#define rf(i,x,y) for(int i=(x);i>=(y);i--)
#define frl(i,x,y) for(int i=(x);i<(y);i++)
using namespace std;
const int N=100002;
const int M=N<<1;
int n,a[N];
int cnt,head[N],Next[M],v[M];

inline void read(int &x){
	char ch=getchar();x=0;
	for(;ch<'0'||ch>'9';ch=getchar());
	for(;ch>='0'&&ch<='9';ch=getchar()) x=(x<<3)+(x<<1)+ch-'0';
}

void add(int x,int y){
	Next[++cnt]=head[x];
	head[x]=cnt;
	v[cnt]=y;
}

int sz[N],f[N],pre[N];
LL tag[N],c[N];
void dfs(int x,int fa){
	sz[x]=1;f[x]=pre[a[x]];
	for(re int i=head[x];i;i=Next[i])
	 if (v[i]!=fa){
	 	pre[a[x]]=v[i];
	 	dfs(v[i],x);
	 	tag[v[i]]+=sz[v[i]];
	 	sz[x]+=sz[v[i]];
	 }
	tag[f[x]]-=sz[x];
	if (f[x]==1) c[a[x]]-=sz[x];
	pre[a[x]]=f[x];
}

LL sum[N];
void dfs(int x,int fa,LL s){
	int cc=c[a[x]];
	s+=tag[x]-cc;
	sum[x]=s;
	for(re int i=head[x];i;i=Next[i])
	 if (v[i]!=fa){
	 	c[a[x]]=tag[v[i]];
	 	dfs(v[i],x,s);
	 }
	c[a[x]]=cc;
}

int b[N];
int main(){
	read(n);
	fr(i,1,n) read(a[i]),b[a[i]]=1;
	int x,y;
	fr(i,2,n){
		read(x);read(y);
		add(x,y);add(y,x);
	}
	//fr(i,1,100000) ver[i].push_back(0);
	int tot=0;
	fr(i,1,100000) if (b[i]) tot++,pre[i]=1,c[i]=n;
	tag[1]=1LL*tot*n;
	dfs(1,0);
	dfs(1,0,0);
	fr(i,1,n) printf("%lld\n",1LL*n*tot-sum[i]);
	return 0;
}