hdu-6035 Colorful Tree

#include <iostream>
#include <vector>
#include <set>
#define LL long long
using namespace std;
int n;
const int N = 200050;
vector<int> g[N];
LL sum[N];//维护以某种颜色为根的子树和
int colors[N];
bool vis[N];
LL wayCnt;
void init()
{
    fill(vis, vis + n + 1, 0);
    fill(sum, sum + n + 1, 0);
    for (int i = 0; i <= n; i++) g[i].clear();
    wayCnt = 0;
}
LL dfs(int now)
{
    int col = colors[now];
    sum[col]++;
    LL preSum = sum[col];//用以计算当前点到某子树之间的联通块大小
    vis[now] = 1;
    LL nowTreeSize = 1;
    int fuck = 0;
    for (int i = 0; i < g[now].size(); i++)
    {
        int nx = g[now][i];
        if (vis[nx]) continue;
        LL siz=dfs(nx);
        LL change = sum[col] - preSum;//计算出底层增量
        LL blockSize = siz - change;
        //cout << now << ": " << blockSize << endl;
        sum[col] += blockSize;
        preSum = sum[col];
        fuck += blockSize;
        wayCnt += blockSize*(blockSize - 1) / 2;
        nowTreeSize += siz;
    }
    return nowTreeSize;
}
int main()
{
    cin.sync_with_stdio(false);
    int cas = 1;
    while (cin >> n)
    {
        set<int> cc;
        for (int i = 1; i <= n; i++)
            cin >> colors[i],cc.insert(colors[i]);
        init();
        for (int i = 0; i < n-1; i++)
        {
            int a, b;
            cin >> a >> b;
            g[a].push_back(b);
            g[b].push_back(a);
        }
        dfs(1);
        //cout << wayCnt << endl;
        for (int i = 1; i <=n; i++)
        {
            if (cc.find(i)!=cc.end()&&i != colors[1])
            {
                LL blockSize = n - sum[i];
                wayCnt += blockSize*(blockSize - 1) / 2;
            }
        }
        cout << "Case #" << cas++ << ": ";
        cout << cc.size()*n*(n-1)/2-wayCnt << endl;
    }
    return 0;
}