hdu 6035(Colorful Tree)

参考题解做的！思路+题意 在代码中藏着呢。

#include <stdio.h>
#include <iostream>
#include <algorithm>
#include <string.h>
#include <math.h>
#include <string>


using namespace std;
typedef long long int LL;
const int MOD= 1e9+7;
const int INF=2e9+1e8;


const int maxn=200000+1000;
/**
 * 数据定义部分
 */
struct Node 
{
    int to,next;
}edge[maxn*2];
int first[maxn],sz,color[maxn];

void addedge(int s,int t)
{
    edge[sz].to=t,edge[sz].next=first[s];
    first[s]=sz++;
}

LL cnt;
LL size[maxn],sum[maxn];
/**
 * sum[i] 是一个维护的数组，表示当前i颜色的此时i颜色为子树的根的子树的大小；
 */
void dfs(int x,int pre)
{
    size[x]=1;
    int all=0;
    int last=sum[color[x]];
    for(int i=first[x];i!=-1;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==pre) continue;
        dfs(to,x);
        size[x]+=size[to];
        LL delta=sum[color[x]]-last;//delta 是变化量
        LL tmp=size[to]-delta;//子树大小-变化量 就是 子树联通块大小；
        last=sum[color[x]];
        cnt+=tmp*(tmp-1)/2;
        all+=size[to]-delta;
    }
    sum[color[x]]+=all+1;
}
/**
 * 题目意思：给定一棵树，树节点编号为1～n，树上的节点有颜色，颜色范围也是1～n。
 * 定义一个值：树上任意不同两个节点路径上的不同颜色的个数（也就是颜色的种类数）。
 * 求所有的任意两个点的这个值的和是多少？有n*(n-1)/2种值，需要加起来。
 * 
 * 大致思路：求不同路径的颜色种类和；
 * 对于每种颜色考虑：有多少条路径含有这种颜色，再把每种颜色的加起来就是
 * 最终答案；
 * 反过来思考颜色：有多少条路径不包含这种颜色，求和，假设求和后值为 sum。
 * 那么最终答案就是：颜色总共的种类数×(n*(n-1)/2)-sum；
 * 那么sum怎么求？我们先只考虑一种颜色，我们考虑的颜色这个点会将树分割开成多个子树，
 * 多个子树内部之间路径是不包含这种颜色的；dfs子树节的时候，遇到的节点颜色就是需要处理的
 * 每个子树的路径数量：假设这个子树有k个节点，那么就是k*(k-1)/2;
 */

 /**
  * 代码具体实现：
  * size[i]代表i号节点的子树的节点数量，
  */

int vis[maxn];
int main()
{
    int cas=1,n;
    while(scanf("%d",&n)+1)
    {
        LL total=0;
        memset(vis,0,sizeof(vis));
        sz=0;
        memset(first,-1,sizeof(first));
        memset(sum,0,sizeof(sum));
        for(int i=1;i<=n;i++) 
        {
            scanf("%d",&color[i]);
            if(vis[color[i]]==0)
            {
                total++;
                vis[color[i]]=1;
            }
        }
        for(int i=1;i<n;i++)
        {
            int x,y;
            scanf("%d%d",&x,&y);
            addedge(x,y);
            addedge(y,x);
        }
        cnt=0;
        dfs(1,1);
        for(int i=1;i<=n;i++)
        {
            if(vis[i]&&i!=color[1])
            {
                LL  temp=n-sum[i];
                cnt+=temp*(temp-1)/2;
            }
        }
        printf("Case #%d: %lld\n",cas++,total*n*(n-1)/2-cnt);
    }
    return 0;
}