HDU 6035 Colorful Tree

版权声明：有错误欢迎大家指出。转载请注明出处~ https://blog.csdn.net/black_miracle/article/details/76474334

Colorful Tree

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

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-1)/2条路都出现  那么答案就是n*(n-1)/2*num

现在去重就是去掉每种颜色没经过多少路径

用一个dfs序搞一搞

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
struct node{
    int to,nex;
}edge[400005];
int head[200005],cnt,num[200005],col[200005],cs,sum[200005],siz[200005],vis[200005];
void add(int u,int v){
    edge[cnt].to=v;
    edge[cnt].nex=head[u];
    head[u]=cnt++;
}
ll ans,de;
void dfs(int u,int pre){  
    siz[u]=1;  
    int step=0;  
    int s1=sum[num[u]];  
    for(int i=head[u];i!=-1;i=edge[i].nex){  
        int v=edge[i].to;  
        if(v!=pre){  
            dfs(v,u);  
            siz[u]+=siz[v];  
            int qr=sum[num[u]]-s1;  
            int szv=siz[v]-qr;  
            s1=sum[num[u]];  
            de=(de+1LL*szv*(szv-1)/2);      
            step+=qr;  
        }  
    }      
    sum[num[u]]+=siz[u]-step;  
}  
int main(){
    int n,i,j,cas=1,u,v;
    while(scanf("%d",&n)!=EOF){
        cs=0;
        de=0;
        cnt=0;
        memset(vis,0,sizeof(vis));
        for(i=1;i<=n;i++){
            head[i]=-1;
            scanf("%d",&num[i]);
            col[++cs]=num[i];
            sum[i]=0;
            siz[i]=0;
            vis[num[i]]=1;
        }
        for(i=1;i<n;i++){
            scanf("%d%d",&u,&v);
            add(u,v);
            add(v,u);
        }
        sort(col+1,col+1+cs);
        cs=unique(col+1,col+1+cs)-col-1;
        ans=(ll)n*(n-1)*cs/2;
        dfs(1,1);
         for(int i=1;i<=n;i++){  
            if(vis[i]&&i!=num[1]){  
                int qr=sum[i];  
                int szv=n-qr;  
                de=(de+1LL*szv*(szv-1)/2);  
            }  
        }
        ans-=de;
        printf("Case #%d: %lld\n",cas++,ans);
    }
    return 0;
}