題目連結

題意：

記 \(i\) 到 \(j\) 的路徑顏色數為 \(s(i,j)\)，對每個 \(i\) 求 \(sum_i=\sum\limits_{j=1}^n s(i,j)\)

\(\text{Solution}\)

考慮在點分治時 \(\text{dp}\)

設當前分治重心為 \(x\) , 正在處理子樹 \(y\) 中全部點此時的貢獻。

對每個顏色考慮 \(x\) 子樹內出除 \(y\) 子樹外所有點與當前點有多少條路徑經過此顏色。

先分析一下 \(\text{dp}\) 時要幹些什麼。

如果一個點 \(u\) 到 \(x\) 之間沒有與 \(x\) 相同顏色的點。

那如果 \(x\)

那 \(u\) 的子樹中的全部點到 \(v\) 的路徑都經過了 \(u\) 的顏色，產生了 \(size(u)\) 的貢獻。

但若 \(v\) 到 \(x\) 的路徑上有與 \(u\) 同色的點，那當前分治對 \(v\) 產生的貢獻就是 \(size(x)-size(y)\) ， \(y\) 是 \(x\) 的一個子節點且 \(v\) 在 \(y\) 的子樹內。

這樣就能知道分治時要如何處理：

先統計出所有 \(x\) 的兒子的子樹中，到 \(x\) 的路徑上沒有與這個點的顏色相同的點的子樹大小之和，對每個顏色 \(a\)

那對於正在處理的 \(x\) 的一個兒子 \(y\) 來說，希望的到的是 \(x\) 的所有除 \(y\) 的子樹的每個顏色的 \(sw\) 之和。

這可以再對 \(y\) 的子樹掃一次，減去每個顏色的 \(sw\) 的一部分。

注意 \(x\) 的顏色的 \(sw\) 需減去 \(size(y)\)。

這樣對 \(y\) 中的子樹內的一個點 \(v\)，就可以很方便地按之前分析的統計了。

時間複雜度僅僅是 \(O(n\log n)\)

程式碼：

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=1e5+10;
int n,m,x,y,edg,nn,rt;char ch;
int to[N<<1],nextn[N<<1],h[N];
int a[N],size[N],mxsz[N],cnt[N];
bool b[N];ll ans[N],sw[N],res,psx;
#define add(x,y) to[++edg]=y,nextn[edg]=h[x],h[x]=edg
inline void read(int &x){
	x=0;ch=getchar();
	while(ch<48||ch>57)ch=getchar();
	while(ch>47&&ch<58)x=(x<<1)+(x<<3)+(ch^48),ch=getchar();
}
void write(ll x){if(x>9)write(x/10);putchar(48+x%10);}
void findrt(int x,int anc){
	int i,y;size[x]=1;mxsz[x]=0;
	for(i=h[x];y=to[i],i;i=nextn[i])if(!b[y]&&y^anc){
		findrt(y,x);
		size[x]+=size[y];
		if(size[y]>mxsz[x])mxsz[x]=size[y];
	}
	mxsz[x]=max(mxsz[x],nn-size[x]);
	if(mxsz[x]<mxsz[rt])rt=x;
}
void dfs(int x,int anc){
	int i,y,ax=a[x];
	++cnt[ax];size[x]=1;
	for(i=h[x];y=to[i],i;i=nextn[i])if(!b[y]&&y^anc){
		dfs(y,x);size[x]+=size[y];
	}
	--cnt[ax];//cnt 都是判斷到x路徑上是否有與其同色的點 
	if(!cnt[ax])sw[ax]+=size[x],res+=size[x];//統計sw及res 
}
void dfs1(int x,int anc){
	int i,y,ax=a[x];
	++cnt[ax];
	for(i=h[x];y=to[i],i;i=nextn[i])if(!b[y]&&y^anc)dfs1(y,x);
	--cnt[ax];
	if(!cnt[ax])sw[ax]-=size[x],res-=size[x];
}
void dfs2(int x,int anc){
	int i,y,ax=a[x];
	++cnt[ax];
	for(i=h[x];y=to[i],i;i=nextn[i])if(!b[y]&&y^anc)dfs2(y,x);
	--cnt[ax];
	if(!cnt[ax])sw[ax]+=size[x],res+=size[x];
}
void dfs_(int x,int anc){
	int i,y,ax=a[x];
	if(!cnt[ax])res+=psx-sw[ax];
	ans[x]+=res;
	++cnt[ax];
	for(i=h[x];y=to[i],i;i=nextn[i])if(!b[y]&&y^anc)dfs_(y,x);
	--cnt[ax];
	if(!cnt[ax])res-=psx-sw[ax];
}
void clear(int x,int anc){
	int i,y;sw[a[x]]=0;
	for(i=h[x];y=to[i],i;i=nextn[i])if(!b[y]&&y^anc)clear(y,x);
}
void work(int x){
	int i,y,ax=a[x];
	res=0;dfs(x,0);
	cnt[ax]=1;
	ans[x]+=res;
	for(i=h[x];y=to[i],i;i=nextn[i])if(!b[y]){
		res-=size[y];
		sw[ax]-=size[y];//這兩行都是減去x的顏色的sw 
		dfs1(y,x);//減去y內的資訊
		psx=size[x]-size[y];
		dfs_(y,x);//統計對y子樹內的貢獻 
		res+=size[y];
		sw[ax]+=size[y];//減完要加回來處理後面的y 
		dfs2(y,x);
	}
	cnt[ax]=0;
	clear(x,0);
}
void solve(int x){
	b[x]=1;
	work(x);
	int i,y;
	for(i=h[x];y=to[i],i;i=nextn[i])if(!b[y]){
		rt=0;nn=size[y];
		findrt(y,x);
		solve(rt);
	}
}
main(){
	read(n);register int i;
	for(i=1;i<=n;++i)read(a[i]);
	for(i=1;i^n;++i)read(x),read(y),add(x,y),add(y,x);
	rt=0;nn=n;mxsz[0]=n;
	findrt(1,0);solve(rt);
	for(i=1;i<=n;++i)write(ans[i]),putchar('\n');
}