HDU6035-Colorful Tree

Colorful Tree

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

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

 

题意：有一棵树，有n个节点，每个节点有个颜色，一条路径的大小为路径上颜色的种类，问所有路径的大小和

解题思路：题意可以转化为所有颜色种类*所有路径数量-每一种颜色有多少路径没有经过，然后dfs就可以O（n）扫出每一种颜色的贡献， Size[k]表示以节点k为根节点的子树的大小，sum[i]表示以颜色i为根节点的子树的大小和，具体看代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <string>
#include <algorithm>
#include <map>
#include <cmath>
#include <set>
#include <stack>
#include <queue>
#include <vector>
#include <bitset>
#include <functional>

using namespace std;

#define LL long long
const int INF = 0x3f3f3f3f;
const LL mod=1e9+7;

int n;
int a[200009],sum[200009];
int Size[200009],vis[200009],x,y;
int s[200009],nt[400009],e[400009];
LL ans;

void dfs(int k,int fa)
{
    Size[k]=1;
    int sum1=0;
    for(int i=s[k];~i;i=nt[i])
    {
        if(e[i]==fa) continue;
        int p=sum[a[k]];
        dfs(e[i],k);
        Size[k]+=Size[e[i]];
        sum1+=sum[a[k]]-p;
        int pp=Size[e[i]]-(sum[a[k]]-p);
        ans-=1LL*(pp-1)*pp/2;
    }
    sum[a[k]]+=Size[k]-sum1;
}

int main()
{
    int cas=0;
    while(~scanf("%d",&n))
    {
        memset(sum,0,sizeof sum);
        memset(s,-1,sizeof s);
        memset(vis,0,sizeof vis);
        int cnt=0,tot=0;
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&a[i]);
            if(!vis[a[i]]) tot++;
            vis[a[i]]=1;
        }
        for(int i=1;i<n;i++)
        {
            scanf("%d%d",&x,&y);
            nt[cnt]=s[x],s[x]=cnt,e[cnt++]=y;
            nt[cnt]=s[y],s[y]=cnt,e[cnt++]=x;
        }
        ans=1LL*n*(n-1)/2*tot;
        dfs(1,0);
        for(int i=1;i<=n;i++)
        {
            if(!vis[i]) continue;
            ans-=1LL*(n-sum[i])*(n-sum[i]-1)/2;
        }
        printf("Case #%d: %lld\n",++cas,ans);
    }
    return 0;
}