[ABC163F]path pass i

LXXI.[ABC163F]path pass i

思路：

反向考虑。我们计算出不包含任何颜色为$i$的节点的路径的数量，再用总路径数一减就行。

则，我们删去所有颜色为$i$的节点，整棵树就会被分成许多连通块。则不经过任何一个颜色为$i$的节点的路径数量，就是$\sum\dfrac{(\text{连通块大小})*(\text{连通块大小}+1)}{2}$。

设$f[i][j]$表示以$i$为根的子树中，删掉所有颜色为$j$的点后，有多少个点与$i$断开联系。再设$sz[i]$表示子树大小。

则$i$节点所在的连通块大小即为$sz[i]-f[i][j]$。

乍一看这状态是$O(n^2)$的。但是如果我们用std::map维护状态，并且在合并状态时启发式合并一下，复杂度就是$O(n\log^2n)$的。

在节点$i$时，将所有$f\big[j\big]\big[col[i]\big]$（其中$j$是$i$的儿子）计入答案，它们被看作是一个连通块的根。

代码：

#include<bits/stdc++.h>
using namespace std;
#define sqr(x) 1ll*x*(x+1)>>1ll
typedef long long ll;
int n,col[200100],sz[200100];
map<int,int>f[200100];
vector<int>v[200100],res[200100];
void dfs(int x,int fa){
	sz[x]=1;
	for(int y:v[x]){
		if(y==fa)continue;
		dfs(y,x),sz[x]+=sz[y];
		int cc=sz[y];
		if(f[y].find(col[x])!=f[y].end())cc-=f[y][col[x]];
		res[col[x]].push_back(cc);
		if(f[x].size()<f[y].size())swap(f[x],f[y]);
		for(auto i:f[y])f[x][i.first]+=i.second; 
	}
	f[x][col[x]]=sz[x];
}
int main(){
	scanf("%d",&n);
	for(int i=1;i<=n;i++)scanf("%d",&col[i]);
	for(int i=1,x,y;i<n;i++)scanf("%d%d",&x,&y),v[x].push_back(y),v[y].push_back(x);
	dfs(1,0);
	for(int i=1;i<=n;i++)res[i].push_back(n-f[1][i]);
	for(int i=1;i<=n;i++){
		ll ans=sqr(n);
		for(auto j:res[i])ans-=sqr(j);
		printf("%lld\n",ans);
	}
	return 0;
}