hdu 6035(容斥)

统计对于每种颜色不含这种颜色的路的数量，用总数剪掉就行。

当我们走到一个节点时，在经历每个子节点时，sum都有一个增量，即离这个节点最近的相同颜色的子节点的子树大小，我们可以利用这个增量计算不包含这个颜色的路径的数量，同时也可以利用这个增量来更新sum。

#include <cstdio>
#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
#include <string>

using namespace std;
const int maxn = 200005;
int c[maxn], sz[maxn], sum[maxn], vis[maxn];
vector<int> G[maxn];
long long ans;
void dfs(int u, int fa)
{
    int color = c[u];
    ++sum[color];
    sz[u] = 1;
    for(int i = 0; i < G[u].size(); ++i)
    {
        int v = G[u][i];
        if(v != fa)
        {
            int prev = sum[color];
            dfs(v, u);
            int part = sum[color] - prev;
            sum[color] += sz[v] - part;
            ans += 1LL * (sz[v] - part) * (sz[v] - part - 1) / 2;
            sz[u] += sz[v];
        }
    }
}
int main()
{
    int n;
    int kase = 0;
    while(~scanf("%d", &n))
    {
        int cnt = 0;
        memset(vis, 0, sizeof(vis));
        memset(sum, 0, sizeof(sum));
        for(int i = 1; i <= n; ++i)
        {
            scanf("%d", &c[i]);
            if(!vis[c[i]])
                ++cnt;
            vis[c[i]] = 1;
            G[i].clear();
            sz[i] = 0;
        }
        for(int i = 0; i < n - 1; ++i)
        {
            int u, v;
            scanf("%d%d", &u, &v);
            G[u].push_back(v);
            G[v].push_back(u);
        }
        ans = 0;
        dfs(1, -1);
//        cout << "ans == " << ans << endl;
        for(int i = 1; i <= n; ++i)
        {
            if(vis[i] && i != c[1])
            {
                ans += 1LL * (n - sum[i]) * (n - sum[i] - 1) / 2;
//                cout << "sum == " << sum[i] << endl;
            }
        }
        printf("Case #%d: %I64d\n", ++kase, 1LL * n * (n - 1) / 2 * cnt - ans);
    }
}