HDU - 6035(71/600)

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

这个题如果按照题意直接做
下辈子也做不出来

首先转化成全部-非法
这种思路
然后去算非法

非法的就是在算颜色a的时候，把所有颜色a的点全都从图上擦掉
剩下的所有路径全都是非法的

但是显然不可能每次都擦掉复杂度爆炸

所以就有了he数组…
这个数组看上去很难理解
但是仔细想想可以理解成一个前缀和
一个存了所有搜过的点的连通块的点数的前缀和..
所以当前搜完子树正在处理的点的he值减掉处理之前的he值的差
代表的意义就是搜当前子树之后全部连通块点的个数-搜这个点之前的全部连通块的点的个数
也就是说这个差是当前子树搜过的连通块的点数
那么用这个子树的节点数减掉已经搜过的连通块的点数
就是靠当前节点最近的连通块的点数了

在他处理完全部这个点以后
再把所有新的连通块加入当前颜色
返回他爹节点

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define int long long 
vector<int>tu[200001];
int ys[200001],he[200001],n,q,w,mp[200001],sh=0,jd=0,zjd[200001],bj[200001],u=0;
void dfs(int gen)
{
    bj[gen]=1;
    zjd[gen]=1;
    int qb=0;
    int sg=he[ys[gen]];
    for(int a=0;a<tu[gen].size();a++)
    {
        int zg=tu[gen][a];
        if(bj[zg])continue;
        dfs(zg);
        int cha=he[ys[gen]]-sg;
        zjd[gen]+=zjd[zg];
        sg=he[ys[gen]];
        int suan=zjd[zg]-cha;
        jd+=suan*(suan-1)/2;
        qb+=suan;
    }
    he[ys[gen]]+=qb+1;
}
main()
{
    while(cin>>n)
    {
        sh=0;
        jd=0;
        for(int a=1;a<=n;a++)tu[a].clear();
        memset(he,0,sizeof(he));
        memset(mp,0,sizeof(mp));
        memset(zjd,0,sizeof(zjd));
        memset(bj,0,sizeof(bj));
        for(int a=1;a<=n;a++)scanf("%lld",&ys[a]);
        for(int a=1;a<=n;a++)
        {
            if(mp[ys[a]])continue;
            mp[ys[a]]=1;
            sh++;
        }
        for(int a=1;a<n;a++)
        {
            scanf("%lld%lld",&q,&w);
            tu[q].push_back(w);
            tu[w].push_back(q);
        }
        ll dan=0;
        dan=sh*n;
        dan*=(n-1);
        dan/=2;
        dfs(1);
        dan-=jd;
        for(int a=1;a<=n;a++)
        {
            if(!mp[a])continue;
    //      if(ys[1]==mp[a])continue;
            dan-=(n-he[a])*(n-he[a]-1)/2;
        }
        printf("Case #%lld: %lld\n",++u,dan);
    }
}