4.19 ABC F path pass i 容斥 树形dp

LINK：path pass i

原本想了一个点分治 yy了半天 发现重复的部分还是很难减掉 况且统计答案的时候有点ex.

(点了别人的提交记录 发现dfs就过了

于是yy了一个容斥 发现可以直接减掉不合法方案。

对于某个点的总方案 ：(1+frac{ncdot (n-1)}{2})

考虑不合法方案 可以发现在树上 我们按顺序便利树 不合法的情况只有两个颜色相同的点之间的那部分的点对不合法。

以及 最后靠上的那部分点的点对是不合法的。

所以 我们统计这些不合法点对的方案即可。

值得注意的是 最后靠上的那部分要注意减掉。

const int MAXN=200010;
int n,len;
ll ans[MAXN];
int a[MAXN],sz[MAXN],s[MAXN];
int lin[MAXN],ver[MAXN<<1],nex[MAXN<<1];
inline void add(int x,int y)
{
	ver[++len]=y;
	nex[len]=lin[x];
	lin[x]=len;
}
inline ll SUM(int x){return (ll)x*(x-1)/2;}
inline void dfs(int x,int father)
{
	sz[x]=1;int ss=s[a[x]],pre=s[a[x]];
	go(x)
	{
		if(tn==father)continue;
		dfs(tn,x);
		sz[x]+=sz[tn];
		ans[a[x]]-=SUM(sz[tn]-(s[a[x]]-pre));
		pre=s[a[x]];
	}
	s[a[x]]=ss+sz[x];
}
int main()
{
	freopen("1.in","r",stdin);
	get(n);
	rep(1,n,i)get(a[i]),++ans[a[i]],ans[i]+=SUM(n);
	rep(2,n,i)
	{
		int get(x);int get(y);
		add(x,y);add(y,x);
	}
	dfs(1,0);
	rep(1,n,i)ans[i]-=SUM(n-s[i]);
	rep(1,n,i)putl(ans[i]);
	return 0;
}