HDU 6035 Colorful Tree（虚树）

Colorful Tree

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

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

31 2 11 22 361 2 1 3 2 11 21 32 42 53 6

 

Sample Output

Case #1: 6Case #2: 29

 

Source

2017 Multi-University Training Contest - Team 1 

 

Recommend

liuyiding

 

题目大意：

    给你一棵n(2<=n<=200000)个节点的每个节点有一个颜色的树，求树上所有路径经过的颜色数的和。

解题思路：

    首先分开考虑每种颜色，也就是拆成n棵只存在当前枚举的颜色，其他颜色视为无色的树。那么就变成了求所有路径中经过这种颜色的路径数，这样还是不好求，再转化一下，如果我们知道所有路径中没有经过这种颜色的路径数也可以算出答案。

   对于求没有进过这个颜色的路径数，我们可以把这棵树上这种颜色的点删掉，就得到了一些联通块，每个联通块中的全部路径都是要求的。于是就可以对每棵树dfs一遍，得到结果，时间复杂度是O(N*N)，显然会TLE。

    继续优化，可以发现，在dfs的过程中，我们只关心和当前节点颜色相同的节点，于是我们就可以考虑使用虚树，也就是不用真的把原树拆开，我们只需要同时维护所有颜色的信息，在dfs到一个节点是只使用和更新当前颜色的信息即可。这样就可一遍dfs解决所有问题，时间复杂度就减为O(N)，十分完美。

AC代码：

#include <iostream>#include <algorithm>#include <cstdio>#include <cstring>#include <cstdlib>#include <cmath>#include <ctime>#include <vector>#include <queue>#include <stack>#include <deque>#include <string>#include <map>#include <set>#include <list>using namespace std;#define INF 0x3f3f3f3f#define LL long long#define fi first#define se second#define mem(a,b) memset((a),(b),sizeof(a))const int MAXN=200000+3;int N, color[MAXN];vector<int> G[MAXN];LL ans;int top[MAXN];//此时颜色为i的联通块的上界int not_in[MAXN];//以i为根节点的子树中，不在i所在的联通块中节点的个数int not_in_root[MAXN];//以颜色i为分界线，不在根节点所在的联通块中的节点个数void init()//初始化{    for(int i=0;i<=N;++i)    {        not_in_root[i]=0;        G[i].clear();    }}inline LL get_path_num(LL num)//得到一个大小为num的联通块中有多少条路径{    return num*(num-1)>>1;}int dfs(int u, int fa){    int now_top=top[color[u]];//以u的颜色为边界的u上面的联通块的上界    top[color[u]]=u;//更新当前颜色的上界    int u_num=1;//以u为根节点的子树的节点数    for(int i=0;i<G[u].size();++i)    {        int v=G[u][i];        if(v==fa)            continue;        not_in[u]=0;        int v_num=dfs(v, u);//以v为根节点的子树的节点数        ans-=get_path_num(v_num-not_in[u]);        u_num+=v_num;    }    (now_top?not_in[now_top]:not_in_root[color[u]])+=u_num;//更新不在当前联通块中的节点数    top[color[u]]=now_top;//回溯    return u_num;}int main(){    int cas=1;    while(~scanf("%d", &N))    {        init();        for(int i=1;i<=N;++i)            scanf("%d",&color[i]);        for(int i=0;i<N-1;++i)        {            int u, v;            scanf("%d%d", &u, &v);            G[u].push_back(v);            G[v].push_back(u);        }        ans=get_path_num(N)*N;//假设所有路径都经过所有颜色        dfs(1, -1);        for(int i=1;i<=N;++i)            ans-=get_path_num(N-not_in_root[i]);//减去根节点所在的联通块        printf("Case #%d: %lld\n",cas++,ans);    }        return 0;}