AtCoder Beginner Contest 163 F path pass i 树的子树节点数+组合数学+容斥原理

AtCoder Beginner Contest 163   比赛人数11511  惨,比赛开始后28分钟看到所有题

AtCoder Beginner Contest 163  F  path pass i   树的子树节点数+组合数学+容斥原理

总目录详见https://blog.csdn.net/mrcrack/article/details/104454762

在线测评地址https://atcoder.jp/contests/abc163/tasks/abc163_f

造了一组数据

对应图形如下

Input:
9
2 2 1 2 2 3 3 3 1
1 2
2 3
3 4
4 5
5 6
6 7
7 8
8 9

Output:
27
34
29
0
0
0
0
0
0

以颜色1为例进行计算
总的数量9*(9+1)/2=45,对应代码ans[i]=get_num(n);

扣除两个颜色1之间无效的,2个颜色2,3个颜色3,对应的非颜色1数量是5,
扣除数量5*(5+1)/2=15,对应代码ans[c]-=get_num(size[v]-dt);

扣除离根节点最近的颜色1到根节点，无效的2个颜色2，扣除数量2*(2+1)/2=3
最后有效数量是45-15-3=27,对应代码ans[i]-=get_num(n-sum[i]);

参考了https://www.cnblogs.com/zdragon1104/p/12734682.html代码

AC代码如下

#include <stdio.h>
#define LL long long
#define maxn 200010
int color[maxn],head[maxn],tot,n,size[maxn],sum[maxn];
LL ans[maxn];
struct node{
	int to,next;
}e[maxn<<1];
void add_edge(int u,int v){
	tot++,e[tot].to=v,e[tot].next=head[u],head[u]=tot;
}
LL get_num(LL x){
	return x*(x+1)/2;
}
void dfs(int u,int fa){
	int v,b,c,t,dt;
	LL pre;
	size[u]=1,c=color[u],pre=sum[c];//非u子树已算出的c颜色为根的节点数量
	for(b=head[u];b;b=e[b].next){
		v=e[b].to;
		if(v==fa)continue;
		t=sum[c];
		dfs(v,u);
		dt=sum[c]-t;//遍历了v之后，c颜色对应的变化量
		size[u]+=size[v];
		ans[c]-=get_num(size[v]-dt);//两个c颜色之间的，其他颜色节点数量
	}
	sum[c]=pre+size[u];//总的c颜色为根的节点数量
}
int main(){
	int i,u,v;
	scanf("%d",&n);
	for(i=1;i<=n;i++)scanf("%d",&color[i]);
	for(i=1;i<n;i++){
		scanf("%d%d",&u,&v);
		add_edge(u,v),add_edge(v,u);
	}
	for(i=1;i<=n;i++)ans[i]=get_num(n);
	dfs(1,0);
	for(i=1;i<=n;i++)ans[i]-=get_num(n-sum[i]);//n-sum[i]指离树根最近的i颜色之外的其他颜色(整棵树扣除所有以i颜色为根的子树)数量
	for(i=1;i<=n;i++)printf("%lld\n",ans[i]);
	return 0;
}