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 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
Source
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;
}