hdu 6035 Colorful Tree

Problem

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

Reference

HDU 6035 树形dp

Meaning

一棵树 n 个结点，颜色为 ci。

定义树上路径的价值为：路径上不同颜色的数量。

求所有路径的的价值总和。

Analysis

正着想，就是对每一种颜色算贡献：有多少条路径进过这种颜色的结点，要去重，有点难做。

反着想，对每一种颜色算有多少条路径不经过这种颜色。

每一种颜色的结点，都把树分割成若干个各
4000
不相连的连通块，连通块内任意两个结点间的路径都不经过该颜色。

于是深搜，对当前搜到的结点 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;
}