CF600E(dsu on tree)

題意:

給你一棵n個點的有根樹,以1爲根,每個點有一個顏色。對於一顆以v爲根的子樹,其出現次數最多的顏色,爲dominating頂點v這個子樹的顏色,可能有多個這樣的顏色。求以n個頂點分別爲根的子樹中,dominating 的顏色和。

思路:

一個很好想的暴力思路就是,dfs n個點,統計該點爲根的子樹中出現的顏色次數。時間、空間都是O(N^2)的。
我們可以每次統計完一個點後清空cnt數組,將空間降到O(N)。

繼續優化,很顯然,某個點的ans應該與其所有孩子的ans有關,所以我們應該儘量將其孩子的信息給用上。

仔細想一下上面暴力的過程,我們會發現其實在清除的時候,最後一棵子樹是沒必要清的,我們可以把它對cntcnt數組的貢獻一直傳上去。於是我們想到了如果人爲地使重兒子是最後一棵子樹的話,豈不美哉。於是我們在dfsdfs時最後訪問重兒子就行了,然後複雜度就被優化到O(nlogn)了。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn = 1e5+5;

int n,col[maxn];
vector<int> G[maxn];
ll ans[maxn];
int sz[maxn],son[maxn],cnt[maxn];
void dfs1(int u,int f){//get son 重孩子
    sz[u]=1;
    for(auto v : G[u]){
        if(v==f) continue;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
}
bool vis[maxn];
int maxv;
ll sum;
//計算貢獻
void calc(int u,int f,int k){//k的取值爲1,-1分別對應累加和清楚

    cnt[col[u]] += k;//把當前節點的顏色累加到cnt中,+1|-1
    if(k>0 && cnt[col[u]]>=maxv){//更新答案
        if(cnt[col[u]] > maxv) sum=0,maxv=cnt[col[u]];
        sum += col[u];
    }
    for(int v : G[u]){
        if(v!=f && !vis[v]) calc(v,u,k);//遞歸計算子節點,vis表示點v已經計算過了
    }
}
void dfs2(int u,int f,int keep){//keep爲1表明當前節點在重兒子的子樹中,需要保留答案
    for(int v : G[u]){
        if(v!=f && v!=son[u]) dfs2(v,u,0);// 先dfs輕兒子
    }
    if(son[u]){      //再dfs重兒子
        dfs2(son[u],u,1);
        vis[son[u]] = 1;//標記重兒子
    }
    calc(u,f,1);//合併輕兒子和重兒子的數據
    ans[u]=sum;
    if(son[u]) vis[son[u]]=0;//取消標記
    if(!keep) calc(u,f,-1),sum=maxv=0;//清楚輕兒子
}
int main(){
    //freopen("in.txt","r",stdin);
    scanf("%d",&n);
    for(int i=1;i<=n;i++) scanf("%d",&col[i]);
    for(int u,v,i=1;i<n;i++){
        scanf("%d%d",&u,&v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs1(1,0);
    dfs2(1,0,0);
    for(int i=1;i<=n;i++) cout<<ans[i]<<' '; puts("");
    return 0;
}

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章