Hdu 6035 Colorful Tree【思维+活用补集】好题~

Colorful Tree

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

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种颜色的点，那么就有ans=n*n*(n-1)/2.

而真实的答案是要在ans的基础上向下减的，那么减少的内容是什么呢？

我们不妨分颜色去考虑，对于每一种颜色，我们ans-不经过这种颜色的路径数即可。

②那么这些个需要减去的路径数的个数要如何统计呢？

如果暴力去做的话，我们可以将所有的颜色分成n棵树，没棵树都只有一种颜色，那么去统计的话，时间复杂度会达到O（n^2）；显然是不行的。

我们考虑优化。

设定size【u】表示以u为根节点，其子树的点的个数，size【v】表示以v为根节点，其子树的点的个数。

设定Sum【color【u】】表示当前Dfs到节点u，以u为根节点的所有子树中距离节点u最近的和颜色u相同的点z的size【z】的和。

那么过程更新维护Sum【color【u】】即可。

每遍历一颗以v为根的子树之后，sum【color【u】】都会增长。

那么此时有temp=size【v】-（Sum【color【u】】-pre），表示从u到以v为根的子树中距离u最近的那个和u颜色相同的点z的路径上点的个数。

那么这些点组成的一个联通块中所有的路径都不会经过颜色color【u】，那么过程ans-=temp*（temp-1）/2即可；

③具体细节参考代码。

Ac代码：

#include<stdio.h>
#include<string.h>
#include<iostream>
#include<vector>
using namespace std;
#define ll long long int
vector<int>mp[250000];
int color[250000];
int size[250000];
int sum[250000];
ll ans;
void Dfs(int u,int from)
{
    ll add=0;
    size[u]=1;
    for(int i=0;i<mp[u].size();i++)
    {
        int v=mp[u][i];
        if(v==from)continue;
        ll pre=sum[color[u]];
        Dfs(v,u);
        size[u]+=size[v];
        ll temp=size[v]-(sum[color[u]]-pre);
        ans-=(temp)*(temp-1)/2;
        add+=temp;
    }
    sum[color[u]]+=add+1;
}
int main()
{
    int kase=0;
    int n;
    while(~scanf("%d",&n))
    {
        memset(sum,0,sizeof(sum));
        memset(color,0,sizeof(color));
        memset(size,0,sizeof(size));
        for(int i=1;i<=n;i++)mp[i].clear();
        for(int i=1;i<=n;i++)scanf("%d",&color[i]);
        for(int i=0;i<n-1;i++)
        {
            int x,y;
            scanf("%d%d",&x,&y);
            mp[x].push_back(y);
            mp[y].push_back(x);
        }
        ans=(ll)(n)*(ll)(n)*(ll)(n-1)/2;
        Dfs(1,-1);
        for(int i=1;i<=n;i++)
        {
            ll temp=n-sum[i];
            ans-=(temp)*(temp-1)/2;
        }
        printf("Case #%d: %lld\n",++kase,ans);
    }
}