HDU6035 Colorful Tree（dfs继承思想，虚树怎么写？？）

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-1)/2(即所有)的路径长度之和。

思路：

看题解才知道怎么做，但是代码怎么实现还是理解了很久很久。
有一个反向思维，如果每一条路径上都含有所有颜色，那么
ANS=n*(n-1)/2 *numcolor
但是实际情况是有的路径上不含这个颜色，所以我们需要把每种颜色没有涉及的路径条数找到，然后用ANS减去就是答案。

朴素：
首先最暴力的方法是，对于每一种颜色我们需要把它从图上面扣掉，然后剩下一些联通块，这些联通块中的路径总条数就可以得到sigma（x*(x-1)/2） （x是联通块中的点数），这样就可以得到不含该颜色的路径条数。直接做是n^2的。
然后队友说虚树就可以把这些取出来O(2)处理，然后总的就变成了O（n）,,,,,,,,emmmmmmm不太会

优化：
我用的就是网上题解中类似树形dp？的方法，可以说就是一个继承的思想。只需要dfs遍历一遍就可以了。
sum[i]记录的是 i这个颜色的截断值，截断值是个什么东西呢，简单来说就是把i这个颜色从图里扣掉 它会带走多少个点（包括它自身）。
这个sum[i]在dfs中是不断更新的，我们维护的核心就是这个值。
size[i]就是在dfs中，以i为根的子树的大小。

理解了sum[i]就很好理解代码了！！！！

最后需要再扫一遍整个数组，因为除了dfs的起点第一个点以外其他的点，他们子树之外的联通块是没有被计算的。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
using namespace std;
const long long MOD=1e9+1;
vector<int> G[200005];
long long ans,ANS;
int sum[200005],size[200005],col[200005],flag[200005];
void dfs(int u,int fa){
    int gone=0;
    size[u]=1;
    for(int i=0;i<G[u].size();i++){
        int to=G[u][i];
        if(to==fa) continue;
        int pre=sum[col[u]];
        dfs(to,u);
        size[u]+=size[to];
        int t=sum[col[u]]-pre;
        gone+=t;
        long long count = size[to]-t;
        ans+=count*(count-1)/2;
    }
    sum[col[u]]+=size[u]-gone; //这里gone是已经加到sum[col[u]]里面的数，所以要减去，避免重复加了。
}
int main(){
    freopen("in.txt","r",stdin);
    int cas=0,n;
    while(~scanf("%d",&n)){

        memset(flag,0,sizeof flag);
        memset(col,0,sizeof col);
        memset(sum,0,sizeof sum);

        for(int i=1;i<=n;i++){
            scanf("%d",&col[i]); 
            flag[col[i]]=1;
        }
        long long gg=0;
        for(int i=1;i<=n;i++){
            if(flag[i]) gg++;
        }
        for(int i=1;i<=n;i++) G[i].clear();
        int x,y;
        for(int i=1;i<=n-1;i++){
            scanf("%d%d",&x,&y);
            G[x].push_back(y);
            G[y].push_back(x);
        }
        ans=0;
        ANS=gg*n*(n-1)/2;
        dfs(1,-1);
        for(int i=1;i<=n;i++){
            if(flag[i]==0 || col[1]==i) continue;
            long long count = n-sum[i];
            ans+=count*(count-1)/2;
        }
        printf("Case #%d: %lld\n",++cas,ANS-ans);
    }
    return 0;
}