HDU_【2017 Multi-University Training Contest 1】——1003 color tree

HDU-6035

题目链接

题目意思

给你一个包含n个节点的树，用一个数字代表一种颜色，树上的路径权值为在这条路径上包含的颜色数量，这棵树总共有n*(n-1)/2条路径。求这棵树上所有路径的总权值。

解题思路

我们要考虑每条路上有多少颜色经过，这样求所有路径上的权值和，但是这种方法不好实现。所以我们反过来思考，我们假设每条路径上都包含了所有的颜色，那么总的路径权值和就是用颜色数*路径数，那么我们现在只要找到每种颜色没有经过那条路径，再用总数减去每种颜色没有经过的路径和就能求出这棵树中所有路径的权值和。

思路实现

所有路径都经过所有颜色，这时候权值可以为颜色的数量乘以总的路径数即 num((n-1)n/2)。

这道题目的难点就是过来求颜色为i的时候，没有经过颜色i的路径之和为多少。

没有经过颜色i的分两种：

1. 在子树中，部分节点k个都没有经过i颜色，那么路径的条数就应该加上(k-1)k/2。该过程在dfs的函数中实现。

2. 再算子树的过程中，都是算的没有经过根节点的路径，这时候就可以借助一个sum数组来求出剩下的所有没有经过的i颜色的路径条数。

在dfs的过程中，如果找到子树中k个点都没有经过i颜色，就应该把这几个点组成的路径条数给算上，然后将其颜色模拟成和i颜色相同（sum[i] += k，模拟过程只是改变sum[i]数组的值，并未改变其本身的颜色）。

这样在将根节点包含进去的时候，就不会重复计算子节点中的颜色不同的路径了。经过dfs最后得到的sum[]就表示，整棵树中，模拟过的i颜色的点的个数有几个，用ct = n-sum[i]表示没有经过颜色i的点的个数有多少个，进而可以推出路径的条数。

代码部分

#include <bits/stdc++.h>
using namespace std;
const int N = 4e5+100;
typedef long long ll;
int c[N],vis[N];///c数组表示第i个节点的颜色ci;vis为标记数组
vector<int> e[N];///e数组用来存储树的信息
ll sum[N],size[N];///sum数组模拟颜色改变后，i颜色在整棵树中个数；size[i]表示的是以i为根的点的个数
ll ans;///ans表示不经过颜色为i的路径总条数
void dfs(int x,int y)///x表示当前节点，y表示其前一个节点
{

    size[x]=1;
    sum[c[x]]++;
    ll pre=sum[c[x]];
    for(int i=0; i<e[x].size(); i++)
    {

        if(e[x][i]==y)
            continue;
        dfs(e[x][i],x);
        size[x]+=size[e[x][i]];/// 以x为根的树的点的总个数，当前的这个点还要加上他的子数上的点
        ll count=size[e[x][i]]-(sum[c[x]]-pre);///count表示与当前节点颜色不同的子节点个数
        ans=ans+(1LL*count*(count-1))/2;
        sum[c[x]]+=count;///将算入路径中的颜色模拟成节点x的颜色。这里注意，如果子树中无x颜色的点的个数为1时，也是构不成路径，但是也要算在sum[x]里面，因为这样的节点在整棵树中也是无法构成路径的
        pre=sum[c[x]];
    }
}

int main()
{
    int n,cas=1;
    while(scanf("%d",&n)!=EOF)
    {
        int num=0;
        ans=0;
        memset(sum,0,sizeof(sum));
        memset(vis,0,sizeof(vis));
        for(int i=1; i<=n; i++)
        {
            e[i].clear();
            scanf("%d",&c[i]);
            if(vis[c[i]]==0)
            {
                vis[c[i]]=1;///标记ci颜色
                num++;
            }
        }
        for(int i=1; i<n; i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            e[u].push_back(v);
            e[v].push_back(u);
        }
        dfs(1,0);
        ll ANS = 1LL*num*((1LL)*n*(n-1))/2;/// 这里要算的就是整棵树中所有颜色都经过每一条路径的所有和
        for(int i=1; i<=n; i++)
        {
            if(vis[i])
            {
                ll ct=n-sum[i];///ct表示模拟过后树中没有颜色i的节点个数
                ans+=ct*(ct-1)/2;/// 这里要算的就是整棵树中没有经过颜色i的路径的个数（除去子树中路径的个数）
            }
        }
        printf("Case #%d: %lld\n", cas++, ANS-ans);
    }
}