HDU60 子树路径问题

Colorful Tree

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

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

 

题意：在一颗树上，有不同颜色的节点，以数字代表，而从一个点到另一个点的权值等于这条路径上经过的颜色数量

题目要求这颗数的所有路径价值总和

要直接去求价值总和的话非常困难，于是我们可以知道对于n个点的树，有n * (n - 1) / 2 条路径，那么它最大化情况下的价值就是 总颜色种类 * n * (n - 1) / 2 ;

算出了最大值后，去算每种颜色不能经过的路径数全部加起来，然后用最大值去减掉不存在的就得出了存在的最大值了

用dfs去跑树，但是由于相同颜色的点会多算边，于是应该对于两个相同的颜色的点之间的路径刨去

#include<stdio.h>
#include<vector>
#include<string.h>
using namespace std;
#define ll long long
#define MAXN 200005

vector<ll>vec[MAXN];
ll col[MAXN],size[MAXN],sum[MAXN],mark[MAXN];
ll ans = 0;

void dfs(ll u,ll v){
	ll all = 0;
	size[u] = 1;
	for(int i = 0; i < vec[u].size();i++){
		ll to;
		to = vec[u][i];
		if(to == v)
			continue;
		ll s = sum[col[u]];  // s先保存当前颜色u的截断值 
		dfs(to,u);
		size[u] += size[to];  
//		printf("size[%lld] += size[%lld]  -----  = %lld\n",u,to,size[u]); 
		ll step = sum[col[u]] - s;  // 儿子里面出现了重复的，记录下来，要剪掉 
//		printf("step = %lld\n",step);
		all += step;		//所有的重复记录，最后从左子树到右子树部分，要剪掉所有重复的 
//		printf("all = %lld\n",all);
		ans += (size[to] - step) * (size[to] - step - 1) / 2;
	}
	sum[col[u]] += size[u] - all;  //减去重复 
//	printf("sum[col[%lld]] += size[%lld] - all (%lld - %lld = %lld) = %lld\n",u,u,size[u],all,size[u] - all,sum[col[u]]);
}

int main(){
	ll n,Case = 1;
	while(scanf("%lld",&n) != EOF){
		memset(mark,0,sizeof(mark));
		memset(size,0,sizeof(size));
		memset(sum,0,sizeof(sum));
		for(int i = 0;i < MAXN;i++){
			vec[i].clear();
		}
		ll totCol = 0;
		ans = 0;
		for(int i = 1;i <= n;i++){
			scanf("%lld",&col[i]);
			totCol += mark[col[i]] ^ 1;
			mark[col[i]] = 1;
		}
		for(int i = 0;i < n - 1;i++){
			ll u,v;
			scanf("%lld %lld",&u,&v);
			vec[u].push_back(v);
			vec[v].push_back(u);
		} 
		dfs(1,-1);
		ll ANS = n * (n - 1) / 2 * totCol;
		for(int i = 1;i <= n;i++){
			if(totCol == col[1] || !mark[i])
				continue;
			ans += (n - sum[i]) * (n - sum[i] - 1) / 2; 
		}
		printf("Case #%d: %lld\n",Case++,ANS - ans);
	}
	return 0;
}