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;
}