【Foreign】采蘑菇 [點分治]

采蘑菇

Time Limit: 20 Sec??Memory Limit: 256 MB

Description

　　

Input

　　

Output

　　

Sample Input

　　5
　　1 2 3 2 3
　　1 2
　　1 3
　　2 4
　　2 5

Sample Output

　　10
　　9
　　12
　　9
　　11

HINT

　　

Main idea

　　詢問從以每個點為起始點時，各條路徑上的顏色種類的和。

Solution

　　我們看到題目，立馬想到了O(n^2)的做法，然后從這個做法研究一下本質，我們確定了可以以點分治作為框架。

　　我們先用點分治來確定一個center(重心)。然后計算跟這個center有關的路徑。設現在要統計的是經過center，對x提供貢獻的路徑。

　　我們先記錄一個記錄Sum[x]表示1~i-1子樹中 顏色x 第一次出現的位置的那個點 的子樹和，然后我們就利用這個Sum來解題。

　　我們顯然可以分兩種情況來討論：

　　（1）統計center->x出現顏色的貢獻：
　　　　顯然，這時候，對于center->x這一段，直接像O(n^2)做法那樣記錄一個color表示到目前為止出現的顏色個數，然后加一下即可。再記錄一個record表示當前可有的貢獻和，一旦出現過一個顏色，那么這個顏色在1~i-1子樹上出現第一次以下的點，對于x就不再提供貢獻了，record減去Sum[這個顏色]，然后這樣深搜往下計算即可。

　　（2）統計center->x沒出現過的顏色的貢獻：
　　　　顯然，對于center->x上沒出現過的顏色，直接往下深搜，一開始為record為(All - Sum[center])，一旦出現了一個顏色，record則減去這個Sum。同樣表示不再提供貢獻即可。

　　我們這樣做就可以求出每個子樹前綴對于其的貢獻了，倒著再做一邊即可求出全部的貢獻。統計x的時候，順便統計一下center。可以滿足效率，成功AC這道題。

Code

#include<iostream>  
#include<algorithm>  
#include<cstdio>  
#include<cstring>  
#include<cstdlib>  
#include<cmath>  
using namespace std;

const int ONE = 600005;
const int INF = 214783640;
const int MOD = 1e9+7;

int n,x,y;
int Val[ONE];
int next[ONE],first[ONE],go[ONE],tot;
int vis[ONE];
int Ans[ONE],Sum[ONE];
int All;


int get()
{ 
        int res,Q=1;    char c;
        while( (c=getchar())<48 || c>57)
        if(c=='-')Q=-1;
        if(Q) res=c-48; 
        while((c=getchar())>=48 && c<=57) 
        res=res*10+c-48; 
        return res*Q; 
}

void Add(int u,int v)
{
        next[++tot]=first[u];    first[u]=tot;    go[tot]=v;
        next[++tot]=first[v];    first[v]=tot;    go[tot]=u;
}

namespace Point
{
        int center;
        int Stack[ONE],top;
        int total,Max,center_vis[ONE];
        int num,V[ONE];
        
        struct power
        {
            int size,maxx;
        }S[ONE];
        
        void Getsize(int u,int father)
        {
            S[u].size=1;
            S[u].maxx=0;
            for(int e=first[u];e;e=next[e])
            {
                int v=go[e];
                if(v==father || center_vis[v]) continue;
                Getsize(v,u);
                S[u].size += S[v].size;
                S[u].maxx = max(S[u].maxx,S[v].size);
            }
        }
             
        void Getcenter(int u,int father,int total)
        {
            S[u].maxx = max(S[u].maxx,total-S[u].size);
            if(S[u].maxx < Max)
            {
                Max = S[u].maxx;
                center = u;
            }
               
            for(int e=first[u];e;e=next[e])
            {
                int v=go[e];
                if(v==father || center_vis[v]) continue;
                Getcenter(v,u,total);
            }
        }
        
        void Ad_sum(int u,int father)
        {
            if(!vis[Val[u]])
            {
                Stack[++top] = Val[u];
                All += S[u].size;    Sum[Val[u]] += S[u].size;
            }
            vis[Val[u]]++;
            for(int e=first[u];e;e=next[e])
            {
                int v=go[e];
                if(v==father || center_vis[v]) continue;
                Ad_sum(v,u);
            }
            vis[Val[u]]--;
        }

        void Calc_in(int u,int father,int center,int Size,int f_time,int record)
        {
            if(!vis[Val[u]]) f_time++, record += Size, record -= Sum[Val[u]];
            Ans[u] += record;    Ans[center]+=f_time;
            Ans[u] += f_time;    vis[Val[u]] ++;
            for(int e=first[u];e;e=next[e])
            {
                int v=go[e];
                if(v==father || center_vis[v]) continue;
                Calc_in(v,u,center,Size,f_time,record);
            }
            vis[Val[u]] --;
        }
        
        void Calc_not(int u,int father,int record)
        {
            if(!vis[Val[u]]) record -= Sum[ Val[u] ];
            Ans[u] += record;    vis[Val[u]] ++;
            for(int e=first[u];e;e=next[e])
            {
                int v=go[e];
                if(v==father || center_vis[v]) continue;
                Calc_not(v,u,record);
            }
            vis[Val[u]] --;
        }
        
        void Dfs(int u)
        {
            Max = n;
            Getsize(u,0);
            Getcenter(u,0,S[u].size);
            Getsize(center,0);
            center_vis[center] = 1;
            
            int num=0; for(int e=first[center];e;e=next[e]) if(!center_vis[go[e]]) V[++num]=go[e];
            
            for(int i=1;i<=num;i++)
            {
                int v=V[i];
                int Size = S[center].size - S[v].size - 1;
                vis[Val[center]] = 1;
                Calc_in(v,center,center, Size,1,All - Sum[Val[center]] + Size);
                vis[Val[center]] = 0;
                Ad_sum(v,center);
            }
            while(top) Sum[Stack[top--]]=0;    All=0;
            
            for(int i=num;i>=1;i--)
            {
                int v=V[i];
                vis[Val[center]] = 1;
                Calc_not(v,center, All-Sum[Val[center]]);
                vis[Val[center]] = 0;
                Ad_sum(v,center);
            }
            
            while(top) Sum[Stack[top--]]=0;    All=0;
            for(int e=first[center];e;e=next[e])
            {
                int v=go[e];
                if(center_vis[v]) continue;
                Dfs(v);
            }
        }
        
}

int main()
{      
        n=get();
        for(int i=1;i<=n;i++)    Val[i]=get();

        for(int i=1;i< n;i++)
        {
            x=get();    y=get();
            Add(x,y);
        }
        
        Point:: Dfs(1);
        for(int i=1;i<=n;i++)
            printf("%d\n",Ans[i]+1);
}

?