HDU_【2017 Multi-University Training Contest 1】——1003 color tree

HDU-6035

题目链接

题目意思

解题思路

思路实现

代码部分

#include <bits/stdc++.h>
using namespace std;
const int N = 4e5+100;
typedef long long ll;
int c
,vis
;///c数组表示第i个节点的颜色ci;vis为标记数组
vector<int> e
;///e数组用来存储树的信息
ll sum
,size
;///sum数组模拟颜色改变后，i颜色在整棵树中个数；size[i]表示的是以i为根的点的个数
ll ans;///ans表示不经过颜色为i的路径总条数
void dfs(int x,int y)///x表示当前节点，y表示其前一个节点
{

size[x]=1;
sum[c[x]]++;
ll pre=sum[c[x]];
for(int i=0; i<e[x].size(); i++)
{

if(e[x][i]==y)
continue;
dfs(e[x][i],x);
size[x]+=size[e[x][i]];/// 以x为根的树的点的总个数，当前的这个点还要加上他的子数上的点
ll count=size[e[x][i]]-(sum[c[x]]-pre);///count表示与当前节点颜色不同的子节点个数
ans=ans+(1LL*count*(count-1))/2;
sum[c[x]]+=count;///将算入路径中的颜色模拟成节点x的颜色。这里注意，如果子树中无x颜色的点的个数为1时，也是构不成路径，但是也要算在sum[x]里面，因为这样的节点在整棵树中也是无法构成路径的
pre=sum[c[x]];
}
}

int main()
{
int n,cas=1;
while(scanf("%d",&n)!=EOF)
{
int num=0;
ans=0;
memset(sum,0,sizeof(sum));
memset(vis,0,sizeof(vis));
for(int i=1; i<=n; i++)
{
e[i].clear();
scanf("%d",&c[i]);
if(vis[c[i]]==0)
{
vis[c[i]]=1;///标记ci颜色
num++;
}
}
for(int i=1; i<n; i++)
{
int u,v;
scanf("%d%d",&u,&v);
e[u].push_back(v);
e[v].push_back(u);
}
dfs(1,0);
ll ANS = 1LL*num*((1LL)*n*(n-1))/2;/// 这里要算的就是整棵树中所有颜色都经过每一条路径的所有和
for(int i=1; i<=n; i++)
{
if(vis[i])
{
ll ct=n-sum[i];///ct表示模拟过后树中没有颜色i的节点个数
ans+=ct*(ct-1)/2;/// 这里要算的就是整棵树中没有经过颜色i的路径的个数（除去子树中路径的个数）
}
}
printf("Case #%d: %lld\n", cas++, ANS-ans);
}
}