HDU6035 Colorful Tree(树形dp)

Colorful Tree

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

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

31 2 11 22 361 2 1 3 2 11 21 32 42 53 6

 

Sample Output

Case #1: 6Case #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 

 

题意:给一棵数,每个点都有一种不同的颜色表示.题目要求求出对于每一条路径,不经过的颜色数的总和.

首先把问题转化成,求每种颜色的贡献,每种颜色的贡献就是经过该种颜色的不同路径数目的和,反向思考一下,就是路径总数减去不经过该种颜色的路径数.

我们考虑某个节点u,它的颜色为c,儿子个数为n,如果我们现在知道了,在节点u的儿子中,m个节点的颜色也为c,他们的儿子个数的总和为s,那么该节点的贡献是不是就是C(n-s,2)

我们需要在一遍dfs中解决这个问题

记录两个数组,一个是son,代表每个节点的孩子总数,一个是sum,记录对于每种颜色,当前已经处理过的该颜色的节点的孩子总数.

#include <cstring>#include <iostream>#include <stdio.h>#include <algorithm>#include <cmath>#include <map>#include<time.h>using namespace std;typedef long long ll;const int MAXN=200000+10;struct edge{    int v,next;}e[MAXN<<1];int tot=0;int head[MAXN];int c[MAXN];ll sum[MAXN];ll son[MAXN];ll res=0;void add(int u,int v){    e[tot].v=v;    e[tot].next=head[u];    head[u]=tot++;}void dfs(int u,int p){    son[u]=1;    ll y=sum[c[u]];    ll x=sum[c[u]];    for(int k=head[u];k!=-1;k=e[k].next){        int v=e[k].v;        if(v!=p){            dfs(v,u);            son[u]+=son[v];            ll temp=son[v]-sum[c[u]]+x;            res+=(temp-1)*temp/2;            x=sum[c[u]];        }    }    sum[c[u]]+=son[u]-(sum[c[u]]-y);}int vis[MAXN];void init(){    memset(head,-1,sizeof(head));    memset(sum,0,sizeof(sum));    memset(vis,0,sizeof(vis));    tot=0;    res=0;}int main(){    int cas=1;    int n;    //freopen("1003.in","r",stdin);    while(scanf("%d",&n)!=EOF){        ll cnt=0;        init();        for(int i=1;i<=n;i++){            scanf("%d",c+i);            if(!vis[c[i]]){                cnt++;                vis[c[i]]=1;            }        }        for(int i=0;i<n-1;i++){            int u,v;            scanf("%d%d",&u,&v);            add(u,v);            add(v,u);        }        dfs(1,0);        ll ans=1LL*n*(n-1)*cnt/2;        for(int i=1;i<=n;i++){            if(vis[i]){                ll temp=n-sum[i];                ans-=temp*(temp-1)/2;            }        }        printf("Case #%d: %lld\n",cas++,ans-res);    }}