hdu 6035 Colorful Tree

Colorful Tree

Time Limit: 6000/3000 MS (Java/Others)    Memory Limit: 131072/131072 K (Java/Others)
Total Submission(s): 182    Accepted Submission(s): 46

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

 

Recommend

liuyiding   |   We have carefully selected several similar problems for you:   6044  6043  6042  6041  6040 

 

Statistic |  Submit |  Discuss |  Note 题意：给你一棵树，树上有n个节点，每个节点都有一种颜色，任意两点间的距离为他们所经过的点的颜色的种类数，求距离和为多少。

思路：现场没想出来，看题解后想了好久终于写了出来。问题等价于算每种颜色有多少条路径经过。又等价于所有情况-每一种颜色有多少条路径没有经过这种颜色。然后这个是可以O（n）跑出来的。跑的时候维护几个数值。size[i]以该节点为根的树的节点数。sum[i]比较难说明，具体看代码领悟。。下面给代码：

//#pragma comment(linker, "/STACK:102400000,102400000") 
#include<iostream>  
#include<cmath>  
#include<queue>  
#include<cstdio>  
#include<queue>  
#include<algorithm>  
#include<cstring>  
#include<string>  
#include<utility>
#include<map>
#include<stack>
#include<vector>
#define maxn 200005
#define inf 0x3f3f3f3f
using namespace std;
typedef long long LL;
const double eps = 1e-8;
const int mod = 1e9 + 7;
int c[maxn], head[maxn], len, sum[maxn], size[maxn], vis[maxn];
LL d;
struct node{
	int v, next;
}p[maxn << 1];
void addedge(int u, int v){
	p[len].v = v;
	p[len].next = head[u];
	head[u] = len++;
}
void dfs(int x, int fa){
	LL pre = sum[c[x]];
	size[x] = 1;
	int add = 0;
	for (int i = head[x]; ~i; i = p[i].next){
		if (p[i].v == fa)
			continue;
		dfs(p[i].v, x);
		size[x] += size[p[i].v];
		LL count = size[p[i].v] - sum[c[x]] + pre;
		pre = sum[c[x]];
		add += count;
		d += count*(count - 1) >> 1;
	}
	sum[c[x]] += add + 1;
}
int main(){
	int n, tcase = 1;
	while (~scanf("%d", &n)){
		memset(head, -1, sizeof(head));
		memset(sum, 0, sizeof(sum));
		memset(vis, 0, sizeof(vis));
		d = len = 0;
		LL number = 0;
		for (int i = 1; i <= n; i++){
			scanf("%d", &c[i]);
			if (!vis[c[i]]){
				vis[c[i]] = 1;
				number++;
			}
		}
		for (int i = 1; i < n; i++){
			int u, v;
			scanf("%d%d", &u, &v);
			addedge(u, v);
			addedge(v, u);
		}
		dfs(1, 0);
		LL ans = (number*(n - 1)*n >> 1) - d;
		for (int i = 1; i <= n; i++){
			if (vis[i] && i != c[1]){
				LL count = n - sum[i];
				ans -= count*(count - 1) >> 1;
			}
		}
		printf("Case #%d: %lld\n", tcase++, ans);
	}
}