HDU 6035 Colorful Tree(樹形DP)

Mr_Treeeee發表於2020-04-06

Colorful Tree

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


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 



 





http://blog.csdn.net/Bahuia/article/details/76141574



這篇講的比較好。




#include <iostream>
#include <stdio.h>
#include <string.h>
using namespace std;
#define  LL long long
const LL N = 2e5+10;

struct node
{
    LL v,nxt;
}len[N<<1];
LL sum[N],size[N];
LL head[N],vis[N],col[N];
LL cnt;
LL lenn;
void add(LL u,LL v)
{
    lenn++;
    len[lenn].v=v;
    len[lenn].nxt=head[u];
    head[u]=lenn;
}
LL ans;
LL dfs(LL u,LL p)
{
    LL allson=0,pre;
    size[u]=1;
    for(LL i=head[u];i!=-1;i=len[i].nxt)
    {
        LL v=len[i].v;
        if(v==p) continue;
        pre=sum[col[u]];
        size[u]+=dfs(v,u);
        LL add=sum[col[u]]-pre;
        LL temp=(size[v]-add)*(size[v]-add-1LL)/2LL;
        ans+=temp;
        allson+=size[v]-add;
    }
    sum[col[u]]+=allson+1;
    return size[u];
}
int main()
{
    LL n;
    LL k=0;
    while(~scanf("%lld",&n))
    {
        cnt=0;
        lenn=0;
        ans=0;
        memset(head,-1,sizeof head);
        memset(vis,0,sizeof vis);
        memset(len,0,sizeof len);
        memset(sum,0,sizeof sum);
        for(LL i=1;i<=n;i++)
        {
            scanf("%lld",&col[i]);
            if(!vis[col[i]])
            {
                cnt++;
                vis[col[i]]=1;
            }
        }
        for(LL i=1;i<=n-1;i++)
        {
            LL u,v;
            scanf("%lld %lld",&u,&v);
            add(u,v);
            add(v,u);
        }
        printf("Case #%lld: ",++k);
        if(cnt==1) //只有一種情況
        {
            printf("%lld\n",(LL)n*(n-1LL)/2LL);
        }
        else
        {
            dfs(1,-1);
            for(LL i=1;i<=n;i++)
            {
                if(!vis[i]) continue;
                ans+=(LL)(n-sum[i])*(n-sum[i]-1LL)/2LL;
            }
            printf("%lld\n",(LL)cnt*n*(n-1LL)/2LL-ans);
        }
    }
}










                                    




相關文章