HDU 6035 树形dp

题目链接：http://acm.hdu.edu.cn/showproblem.php?pid=6035
一棵n个结点的树，每个结点都有颜色，定义两点之间的路径长度为路径上出现的不同颜色数目，求树上所有路径的长度和。

https://blog.csdn.net/Bahuia/article/details/76141574原链接

大神的写法简单易懂

#include<iostream>
#include<cstring>
#include<vector>
#include<algorithm>
using namespace std;
const int maxn = 3e5 + 111;
typedef long long ll;
ll sum[maxn];
ll list[maxn];
vector<int>G[maxn];
void add(int be, int en) {
	G[be].push_back(en);
}
ll ans = 0;
ll siz[maxn];
int vis[maxn];

int dfs(int x, int fa) {
	ll cns = 0;
	siz[x] = 1;
	for (int i = 0; i < G[x].size(); i++) {
		int p = G[x][i];
		if (p == fa) continue;
		ll last = sum[list[x]];
		dfs(p, x);//一个儿子一个儿子的算
		siz[x] += siz[p];
		ll add = sum[list[x]] - last;
		ans += 1LL*(siz[p] - add)*(siz[p] - add - 1) / 2;
		cns += siz[p] - add;
	}

	sum[list[x]] += cns + 1;
	return 0;
}



int main() {
	ll n;
	int c = 0;
	while (~scanf("%lld", &n)) {
		memset(sum, 0, sizeof(sum));
		memset(siz, 0, sizeof(siz));
		memset(vis, 0, sizeof(vis));

		ll cnt = 0;
		for (int i = 1; i <= n; i++) {
			scanf("%d", &list[i]);
			if (vis[list[i]] == 0) {
				vis[list[i]] = 1;
				cnt++;
			}
		}

		ans = 0;
		for (int i = 1; i < n; i++) {
			int be, en;
			scanf("%d %d", &be, &en);
			add(be, en);
			add(en, be);
		}


		if (cnt == 1) {//只有一种颜色
			printf("Case #%d: %lld
", ++c, 1LL*n*(n-1)/2);
			for (int i = 0; i <= n; i++) G[i].clear();
			continue;
		}
		dfs(1, -1);
		cnt *= 1LL*n * (n - 1) / 2;//总路径数量
		
		for (int i = 1; i <= n; i++) {
			if (vis[i]) {
				ans += 1LL*(n - sum[i])*(n - sum[i] - 1) / 2;//没选过的儿子们，就算选过也没关系
			}
		}

		for (int i = 0; i <= n; i++) G[i].clear();

		printf("Case #%d: %lld
", ++c, cnt - ans);
	}
	return 0;
}