hdu 6035 Colorful Tree

Problem

acm.hdu.edu.cn/showproblem.php?pid=6035

Reference

HDU 6035 树形dp

Meaning

一棵树 n 个结点，颜色为 ci。
定义树上路径的价值为：路径上不同颜色的数量。
求所有路径的的价值总和。

Analysis

正着想，就是对每一种颜色算贡献：有多少条路径进过这种颜色的结点，要去重，有点难做。
反着想，对每一种颜色算有多少条路径不经过这种颜色。
每一种颜色的结点，都把树分割成若干个各不相连的连通块，连通块内任意两个结点间的路径都不经过该颜色。
于是深搜，对当前搜到的结点 i（颜色为 ci）与在以它为根的子树内，跟它最近的、颜色也是 ci 的后代结点一起，割出了一个颜色不是 ci 的结点组成的连通块，大小为：树 i 的大小减去上述那些后代结点为根的子树的大小之（下图：size[1] - (size[4] + size[7] + size[8])）。

这样 dfs 算出来只是“子树内”的不经颜色 ci 的的路径数，漏了“子树上方”的不经颜色 ci 的路径，如图：

在 dfs 完之后，要补充这些连通块的贡献。

Code

#include <cstdio>
#include <cstring>
#include <vector>
using namespace std;
const int N = 200000;

int c[N+1]; // colour
int sz[N+1]; // i为根的树的大小
int sum[N+1]; // 到目前为止，已经搜过的树根颜色为i的树总大小
bool exist[N+1]; // 颜色i是否存在
vector<int> g[N+1];

int dfs(int now, int fa, long long &ans)
{
    sz[now] = 1;
    int cnt = 0;
    for(int i = 0, e = g[now].size(); i < e; ++i)
    {
        if(g[now][i] == fa)
            continue;
        int last = sum[c[now]], son = g[now][i];
        sz[now] += dfs(son, now, ans);
        long long add = sum[c[now]] - last;
        ans += (sz[son] - add) * (sz[son] - add - 1) >> 1;
        cnt += sz[son] - add;
    }
    sum[c[now]] += cnt + 1;
    return sz[now];
}

int main()
{
    int n;
    for(int kase = 1; ~scanf("%d", &n); ++kase)
    {
        int cnt = 0; // 颜色数
        memset(exist, false, sizeof exist);
        for(int i = 1; i <= n; ++i)
        {
            scanf("%d", c+i);
            if(!exist[c[i]])
            {
                exist[c[i]] = true;
                ++cnt;
            }
        }
        for(int i = 1, x, y; i < n; ++i)
        {
            scanf("%d%d", &x, &y);
            g[x].push_back(y);
            g[y].push_back(x);
        }
        long long ans = 0;
        memset(sz, 0, sizeof sz);
        memset(sum, 0, sizeof sum);
        dfs(1, -1, ans);
        // 补充“上方”连通块
        for(int i = 1; i <= n; ++i)
            if(exist[i])
                ans += (long long)(n - sum[i]) * (n - sum[i] - 1) >> 1;
        ans = (long long)n * (n - 1) * cnt / 2 - ans;
        printf("Case #%d: %I64d\n", kase, ans);
        for(int i = 1; i <= n; ++i)
            g[i].clear();
    }
    return 0;
}