2017ACM暑期多校联合训练 - Team 1 1003 HDU 6035 Colorful Tree （dfs）

题目链接

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

分析：

给你一棵树 n个点，n个点有自己的颜色，每两个点的权值就是这条路上不同的颜色的种类数，问整棵树的权值。。。首先可以想到的。。一棵树上所有的边是 n×(n-1)/2 假如每个颜色都在所有的边上，那么总权值就是 总颜色 num×n×(n-1)/2 那么很简单的一个想法 是

总权值 减去多加上的权值，那么多加上的权值就是每个点的颜色不在的边的数量的总和，由于时间要求，你可能需要在on 内得到这些东西，所以你可以用类似树形dp维护，那么如何得到这个值呢。画个树的图，可以想到，每个树节点下面支路都是相连的，每一部分的求和都可以用 n×（n-1）/2。然后再运用虚树的思想，每个点管理下面它需要管理的部分。。（即只有它自身的这种颜色，别的颜色看作无色，那么每个节点的管理部分都是他下面的点），你可能会想如果每次都处理的是自己下面的点，那么隔壁树枝的怎么办呢，由于这是一个分治合并的过程，在树上那么总会有一个汇点处理这一切的。。。

那么现在要找的路径数 就是 把这个颜色的点都去掉，得到了很多的块或者点，每个块内部的路径可以用n×（n-1)/2 得到，然而每个部分都有相应管辖的点.

#include <bits/stdc++.h>
using namespace std;
const int N = 4e5+100;
typedef long long ll;
int c[N],vis[N];///c[i]表示第i个点的颜色，vis[i]表示的是i这个颜色有没有出现过
vector<int> e[N];///存储这个图呢
ll sum[N],size[N];///sum[i]表示的是这个颜色i所能管理的点，size[i]表示的是以i为根的点的个数
ll ans;
void dfs(int x,int y)
{

    size[x]=1;/// 指的是x为根的树的点的个数
    sum[c[x]]++;/// sum是指 当前颜色的所能管理的点（包括自身的总个数） 那么每种颜色 最后都是一个接近根的值
    ///到不了根的部分 就要在后面剔除
    ll pre=sum[c[x]];/// 指从根到 x ，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);/// 总的减去被管理的部分 剩下的都是自由的联通块 直接求路径后面剔除
        ans=ans+(1LL*count*(count-1))/2;
        //printf("%d\n",sum[c[x]]);
        sum[c[x]]+=count;/// 加上这些被剔除的，相当于把这些与根节点颜色不一样的点标记成颜色一样的，这样递归往上返回的时候就不会重复计算这些点了
        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]])
            {
                vis[c[i]]=1;///标记这个颜色有没有出现过
                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];
                ans+=ct*(ct-1)/2;
            }
        }
        printf("Case #%d: %lld\n", cas++, ANS-ans);
    }
}