Colorful Tree hdu 6035 （虚树，树形dp）

There is a tree with nn nodes, each of which has a type of color represented by an integer, where the color of node ii is cici.

The path between each two different nodes is unique, of which we define the value as the number of different colors appearing in it.

Calculate the sum of values of all paths on the tree that has n(n−1)2n(n−1)2 paths in total.
Input
The input contains multiple test cases.

For each test case, the first line contains one positive integers nn, indicating the number of node. (2≤n≤200000)(2≤n≤200000)

Next line contains nn integers where the ii-th integer represents cici, the color of node ii. (1≤ci≤n)(1≤ci≤n)

Each of the next n−1n−1 lines contains two positive integers x,yx,y (1≤x,y≤n,x≠y)(1≤x,y≤n,x≠y), meaning an edge between node xx and node yy.

It is guaranteed that these edges form a tree.
Output
For each test case, output ” Case #xx: yy” in one line (without quotes), where xx indicates the case number starting from 11 and yy denotes the answer of corresponding case.
Sample Input
3
1 2 1
1 2
2 3
6
1 2 1 3 2 1
1 2
1 3
2 4
2 5
3 6
Sample Output
Case #1: 6
Case #2: 29

给你一棵树 n个点，n个点有自己的颜色，每两个点的权值就是这条路上不同的颜色的种类数，问整棵树的权值。。。首先可以想到的。。一棵树上所有的边是 n*(n-1) 假如每个颜色都在所有的边上，那么总权值就是 总颜色 *n(n-1) 那么很简单的一个想法 是
总权值 减去多加上的权值，那么多加上的权值就是每个点的颜色不在的边的数量的总和，由于时间要求，你可能需要在on 内得到这些东西，所以你可以用类似树形dp维护，那么如何得到这个值呢。画个树的图，可以想到，每个树节点下面支路都是相连的，每一部分的求和都可以用 n×（n-1）/2。然后再运用虚树的思想，每个点管理下面它需要管理的部分。。（即只有它自身的这种颜色，别的颜色看作无色，那么每个节点的管理部分都是他下面的点），你可能会想如果每次都处理的是自己下面的点，那么隔壁树枝的怎么办呢，由于这是一个分治合并的过程，在树上那么总会有一个汇点处理这一切的。。。
那么现在要找的路径数 就是 把这个颜色的点都去掉，得到了很多的块或者点，每个块内部的路径可以用n×（n-1)/2 得到，然而每个部分都有相应管辖的点

#include <bits/stdc++.h>
using namespace std;
const int N = 4e5+100;
typedef long long ll;
int c[N];
vector<int> e[N];
ll sum[N],size[N];
ll ans;
void dfs(int x,int y)
{
    size[x]=1;// 指的是x为根的树的点的个数
    sum[c[x]]++;// sum是指 当前颜色的所能管理的点（包括自身的总个数） 那么每种颜色 最后都是一个接近根的值
    //到不了根的部分 就要在后面剔除
    ll pre=sum[c[x]];// 指从根到 x 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);// 总的减去被管理的部分 剩下的都是自由的联通块 直接求路径后面剔除
        ans=ans+(1LL*count*(count-1))/2;
        sum[c[x]]+=count;// 加上这些被剔除的，因为这些被剔除被合并到当前点被管理了
        pre=sum[c[x]];//更新这个颜色的所有的管理点
    }
}

int vis[N];

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]]) 
        {
            vis[c[i]]=1;
            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];
        ans+=ct*(ct-1)/2;
        }
    }

    printf("Case #%d: %lld\n", cas++, ANS-ans);
    }
}