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);
    }
}