HDU 6035 树形DP

题目链接

题意：

一棵n个结点的树，每个结点都有颜色，定义两点之间的路径长度为路径上出现的不同颜色数目，求树上所有路径的长度和。

思路：

此题以贡献来考虑计数，先假设每一种颜色在树中每一条路径里都贡献了，那么就有答案总贡献N*(N-1)/2*cnt（cnt表示出现的颜色种类数）然后再依次剔除颜色i所不在的路径记录在另一变量ans。剔除过程通过树形DP完成。

剔除过程如下，每次我们遍历到某个根u时去它的子树v中找所有的颜色相同且处在最高位置的点，用sizes[v]减去以这些点为根的子树大小即可获得一块连通块的大小，这个连通块内任意两点的路径根u的颜色是无法贡献的，由此ans += tmp*(tmp-1)/2

定义sum[i]表示遍历到目前点为止，与颜色i相同且处在最高位置其子树和。维护这个数据很简单，初始时初始化0，当遍历完了所有子树后得到一个变量为各子树中连通块大小的和allson，这些连通块分散且不包含根u的颜色，但是把节点u加入这些连通块里他们就会连起来而且包含根u的颜色，将这一大连通块的大小维护进sum[i]就可以实现sum[i]计数的目的了。

C++代码：

#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
const int maxn = 200010;
const int maxm = 400010;

int n,tol,head[maxn];
struct edge
{
    int to,next;
}es[maxm];

void addedge( int u , int v )
{
    es[tol].to = v;
    es[tol].next = head[u];
    head[u] = tol++;
}

int color[maxn],vis[maxn];
LL ans,sum[maxn],sizes[maxn];

void dfs( int u , int f )
{
    sizes[u] = 1;
    LL allson = 0;
    for ( int i=head[u] ; i!=-1 ; i=es[i].next )
    {
        int v = es[i].to;
        if ( v!=f )
        {
            LL last = sum[color[u]];
            dfs( v , u );
            sizes[u] += sizes[v];
            LL add  = sum[color[u]]-last;
            ans += ( sizes[v]-add )*( sizes[v]-add-1 )/2;
            allson += sizes[v]-add;
        }
    }
    sum[color[u]] += allson+1;
}

int main()
{
    for( int cas=1 ; scanf ( "%d" , &n )==1 ; cas++ )
    {
        memset ( vis , 0 , sizeof(vis) );
        memset ( sum , 0 , sizeof(sum) );
        LL cnt = 0,N = n;
        for ( int i=1 ; i<=n ; i++ )
        {
            scanf ( "%d" , &color[i] );
            if ( !vis[color[i]] )cnt++;
            vis[color[i]] = 1;
        }
        tol = 0;
        memset ( head , -1 , sizeof(head) );
        for ( int i=1 ; i< n ; i++ )
        {
            int u,v;
            scanf ( "%d%d" , &u , &v );
            addedge ( u , v );
            addedge ( v , u );
        }
        printf ( "Case #%d: " , cas );
        if ( cnt==1 )
        {
            printf ( "%lld\n" , N*(N-1)/2 );
            continue;
        }
        ans = 0;
        dfs ( 1 , 0 );
        for ( int i=1 ; i<=n ; i++ )
        {
            if ( !vis[i] ) continue;
            ans += ( N-sum[i] )*( N-sum[i]-1 )/2;
        }
        printf ( "%lld\n" , N*(N-1)/2*cnt-ans );
    }
    return 0;
}