首页 > > 程序设计 > C/C++ >

HDU 6035---Colorful Tree（树形DP）

2018-06-17 22:09:38来源：未知阅读 ()

题目链接

Problem Description

There is a tree with

Input

The input contains multiple test cases.

For each test case, the first line contains one positive integers

Output

For each test case, output "Case #" in one line (without quotes), where

Sample Input

1 2 1

1 2

2 3

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个节点的树，编号从 1 到 n ，每个节点涂有一种颜色ci （1<=ci<=n），对于树上的每条路径有个值，表示这条路径上的颜色种类数，现在对树上所有的路径值求和。

思路：可以转换一下题目描述，即求包含颜色ci的路径数，然后对所有颜色求和即可。这样求依然不好计算，可以反向考虑，包含颜色ci的路径数=n*(n-1)/2 - 不包含颜色ci的路径数，求不包含颜色ci的路径数即为颜色ci将整棵树分成一个个的联通块，对每个联通块，如果包含节点数为x,即为C(x,2)，对所有联通块求和；

代码如下：

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <vector>
using namespace std;
typedef long long LL;
const int N=200005;
int c[N],sum[N],sz[N];
vector<int> t[N];
LL ans;

int dfs(int u,int pa)
{
    sz[u]=1;
    int cn=t[u].size();
    for(int i=0;i<cn;i++)
    {
        int v=t[u][i];
        if(v==pa) continue;
        int r=sum[c[u]];
        sz[u]+=dfs(v,u);
        int o=sum[c[u]]-r;
        ans+=(LL)(sz[v]-o)*(LL)(sz[v]-o-1)/2;
        sum[c[u]]+=sz[v]-o;
    }
    sum[c[u]]++;
    return sz[u];
}

int main()
{
    ///cout << "Hello world!" << endl;
    int n,Case=1;
    while(scanf("%d",&n)!=EOF)
    {
        memset(sum,0,sizeof(sum));
        for(int i=1;i<=n;i++){
            scanf("%d",&c[i]);
            t[i].clear();
        }
        for(int i=1;i<n;i++)
        {
            int u,v;  scanf("%d%d",&u,&v);
            t[u].push_back(v);
            t[v].push_back(u);
        }
        ans=0;
        dfs(1,-1);
        for(int i=1;i<=n;i++)
        {
            ans+=(LL)(n-sum[i])*(LL)(n-sum[i]-1)/2;
        }
        ans=(LL)n*(LL)(n-1)*(LL)n/2-ans;
        printf("Case #%d: %lld\n",Case++,ans);
    }
    return 0;
}