HDU 6035 Colorful Tree

【题目链接】http://acm.hdu.edu.cn/showproblem.php?pid=6035

题目意思

给你n个点，每个点都会用一个数字代表一种颜色，之后给你（n-1）相连的条路，问你n总共形成的n*（n-1)/2条路每路经过的颜色的个数的总和（1到2,2到3,1到3是3条路。所以n个点能形成n*(n-1)/2条路）

解题思路

首先题目意思可以转化成每种颜色经过的路的条数总和，有可以转化成全部的路数*颜色个数-每种颜色没有经过的路个数总和。
每种颜色不经过的路数==每种颜色不经的点数*（不经过的点数-1）/2（count*(count-1)/2）；
这样我们就要求每种颜色不经过的点的个数，而每种颜色不经过点的个数可以分为两种：
第一种为与根节点有关联但不颜色不同的点：
如下图以节点2为根节点，节点4就是与2相关联但 不相同的点。节点1为根节点，节点2，4， 5。
第二中为与根节点无关的且颜色不同的点：
如下图以节点2为根节点，节点1，3，6。节点1为根节点，就没有无关的点，节点4为根节点， 剩下的点就全为无关的点了。

为什么要分两种呢？因为无关的点注定不会经过根节点所以组成的无根节点颜色的路数就为：n*（n-1）/2。而相关联的点就要考虑是否在根节点的同一子树上了，不同子树上的点能不能组成路，因为这样注定要经过根节点如下图：对于根节点而言2，3，4，5都是不同颜色，但是组成的不经过红色的路只有2条2-4和3-5，而不是4*3/2=6；

代码部分

#include <bits/stdc++.h>
using namespace std;
const int N = 4e5+100;
typedef long long ll;
int c[N],vis[N];///c数组表示第i个节点的颜色ci;vis为标记数组
vector<int> e[N];///e数组用来存储树的信息
ll sum[N],size[N];///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);
    }
}