HDU 6035 (虚树)(统计颜色)

HDU 6035 Colorful Tree

Problem : 给一棵树，每个结点有一种颜色，定义每条路径的权值为这条路径上颜色的种数，询问所有路径(C(n,2)条)的权值之和。
Solution : 分开考虑每种颜色对答案的贡献。对于一种颜色，一共有C(n,2)条路径，再考虑问题的反面。不含有这种颜色的路径的数量，即为将这棵树上所有这种颜色的点删去后，所有剩下的树的路径数量。
问题关键在于统计删去点之后每颗树的大小，对于每个节点要计算删去其所有相同颜色的子树，用栈进行维护。即每次访问到每个节点时，将其最近的相同颜色的祖先的大小减去该子树的大小。

#include <bits/stdc++.h>

using namespace std;

const int N = 200008;

vector <int> eg[N];

int cl[N];
int flag[N];
int size[N];
int nt[N], st[N];
long long tag[N];
long long tagrt[N];

int n;

void dfs(int u, int fa) 
{
	size[u] = 1;
	for (auto v : eg[u])
	{
		if (v == fa) continue;
		dfs(v, u);
		size[u] += size[v];
	}
}
void solve(int u, int fa)
{
	int rt;
	int last = st[cl[u]];
	if (last == 0) 
	{
		rt = 1; 
		tagrt[cl[u]] += size[u];
	}
	else 
	{
		rt = nt[last];
		tag[rt] += size[u];
	}
	st[cl[u]] = u;
	for (auto v: eg[u])
	{
		if (v == fa) continue;
		nt[u] = v;
		solve(v, u);
	}
	st[cl[u]] = last;
}

int main()
{
	cin.sync_with_stdio(0);
	int cas = 0;
	while (cin >> n)
	{
		for (int i = 1; i <= n; ++i) 
		{
			size[i] = nt[i] = tag[i] = tagrt[i] = flag[i] = 0, eg[i].clear();
		}
		for (int i = 1; i <= n; ++i) cin >> cl[i], flag[cl[i]] = 1;
		int num = 0;
		for (int i = 1; i <= n; ++i) num += flag[i];
		long long ans = 1ll * n * (n - 1) / 2 * num;
		for (int i = 1; i < n; ++i)
		{
			int u, v;
			cin >> u >> v;
			eg[u].push_back(v);
			eg[v].push_back(u);
		}
		dfs(1, 0);
		solve(1, 0);

		for (int i = 2; i <= n; ++i)
		{
			int p = size[i] - tag[i];
			ans -= 1ll * p * (p - 1) / 2;
		}
		for (int i = 1; i <= n; ++i)
			if (flag[i]) 
			{
				int p = size[1] - tagrt[i];
				ans -= 1ll * p * (p - 1) / 2;
			}
		cout << "Case #" << ++cas << ": " << ans << endl;
	}
}