Colorful Tree（HDU 6035）

Colorful Tree

Time Limit: 6000/3000 MS (Java/Others)    Memory Limit: 131072/131072 K (Java/Others)

Total Submission(s): 1844    Accepted Submission(s): 771

Problem Description

There is a tree with n nodes, each of which has a type of color represented by an integer, where the color of node i is ci.

The path between each two different nodes is unique, of which we define the value as the number of different colors appearing in it.

Calculate the sum of values of all paths on the tree that has n(n−1)2 paths in total.

 

Input

The input contains multiple test cases.

For each test case, the first line contains one positive integers n, indicating the number of node. (2≤n≤200000)

Next line contains n integers where the i-th integer represents ci, the color of node i. (1≤ci≤n)

Each of the next n−1 lines contains two positive integers x,y (1≤x,y≤n,x≠y), meaning an edge between node x and node y.

It is guaranteed that these edges form a tree.

 

Output

For each test case, output "Case #x: y" in one line (without quotes), where x indicates the case number starting from 1 and y denotes the answer of corresponding case.

 

Sample Input

3 1 2 1 1 2 2 3 6 1 2 1 3 2 1 1 2 1 3 2 4 2 5 3 6

 

Sample Output

Case #1: 6 Case #2: 29

 

Source

2017 Multi-University Training Contest - Team 1

 

//题意：一棵树里每个节点都有一种颜色，求每个节点到其他所有点的路径里经过的颜色总和，再把这些总和都加起来，就是输出的答案...解释的不是很清楚，题目不长，最好自己看下...

//思路：Answer = 所有颜色种类 * 所有路径数量 - 每一种颜色有多少路径没有经过 .

一开始假设每条路径都经过了所有颜色，再把每种颜色没经过的路径数减去就行，这个应该很好理解。问题是怎么算没经过的路径数？操作是这样的，如果算颜色1没经过的路径数，我们先把图里所有颜色是1的节点遮起来（假设这个点不存在，图在这个点是断路），图就被分成了很多块，每块的值= 那一块里的顶点数*（那一块里的顶点数-1）/2。

所有块的值加起来就是不经过颜色1的所有路径数。

到这里是不是还是很好理解，那么问题来了，怎么实现？...题解里说用虚树什么的...

用一个DFS即可，复杂度O（n）

用Size数组储存以每个节点为根节点的子树大小（即子树里节点的个数），Sum数组...很难解释，大概是表示以每种颜色为根节点的子树的大小和，但不是非常准确，如果以颜色2为根节点的子树里还含有颜色为2的节点，那只要算根节点这个颜色为2的子树大小即可，若在以这个颜色为2的点为根节点的子树之外还有颜色为2的点，那还要加上这个点的值...不知道能不能理解...解释不清楚，大概就这个意思...

以下图颜色2为例，代码里最后的for里减去的（n-sum[2]）*（n-sum[2]-1）/2的那部分减去的是下图橙色圈里的那块，dfs里减去pp那部分是下图里蓝色圈的那块。其他具体的按照自己的理解再思考思考。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <string>
#include <cmath>
#include <vector>
#include <algorithm>
using namespace std;

const int MAX = 200000 + 100;

int n;
int color[MAX];
int vis[MAX];
long long ans;
vector<int>map[MAX];
int Size[MAX];
int sum[MAX];

void dfs(int root, int front)
{
	int sum1 = 0;
	for (int i = 0; i < map[root].size(); i++)
	{
		//防止因为是无向边而造成的死循环
		if (map[root][i] == front)
			continue;

		int p = sum[color[root]];
		dfs(map[root][i], root);
		Size[root] += Size[map[root][i]];
		sum1 += sum[color[root]] - p;

		//pp减掉的是以root的孩子节点为根节点的子树里的不含root的情况
		int pp = Size[map[root][i]] - (sum[color[root]] - p);
		ans -= 1LL * pp*(pp - 1) / 2;
	}
	sum[color[root]] += Size[root] - sum1;
}

int main()
{
	int Case = 1;
	while (scanf("%d", &n) != EOF)
	{
		int x, y;

		//tot表示总共有多少种颜色
		int tot = 0;

		memset(vis, 0, sizeof(vis));
		memset(sum, 0, sizeof(sum));
		for (int i = 0; i <= n; i++)
			map[i].clear();
		for (int i = 1; i <= n; i++)
		{
			scanf("%d", &color[i]);
			Size[i] = 1;
			if (!vis[color[i]])
			{
				vis[color[i]] = 1;
				tot++;
			}
		}
		for (int i = 1; i <= n - 1; i++)
		{
			scanf("%d%d", &x, &y);

			//注意是无向边，x能到y，y也能到x
			map[x].push_back(y);
			map[y].push_back(x);

		}
		ans = 1LL * n*(n - 1)*tot / 2;
		dfs(1, 0);
		printf("Case #%d: ", Case++);
		for (int i = 1; i <= n; i++)
		{
			if (!vis[i])
				continue;
			ans -= 1LL * (n - sum[i]) * (n - sum[i] - 1) / 2;
		}
		printf("%lld\n", ans);
	}
	return 0;
}