1. 程式人生 > 實用技巧 >樹上啟發式合併/dsu on tree

樹上啟發式合併/dsu on tree

樹上啟發式合併/dsu on tree

前置芝士

啟發式合併和樹鏈剖分的部分知識。(不會的去這裡搜)

因為要在一顆樹上進行啟發式合併,所以要找最優的方法,即優雅的暴力(霧

它可以讓 \(O(n^2)\) 變為 \(O(n\log n)\) (證明 你就想想啟發式合併就完了)

概念

樹上啟發式合併(dsu on tree)對於某些樹上離線問題可以速度大於等於大部分演算法且更易於理解和實現的演算法。

因為沒什麼需要講的,所以直接看題

例題1

U41492 樹上數顏色
題意

給一棵根為1的樹,每次詢問子樹顏色種類數

思路

​ 直接暴力預處理的時間複雜度為 O(n^2) ,即對每一個子節點進行一次遍歷,每次遍歷的複雜度顯然與 n 同階,有 n 個節點,故複雜度為 O(n^2)

​ 可以發現,每個節點的答案是其子樹的疊加,考慮利用這個性質處理問題

​ 我們可以先預處理出每個節點子樹的size和它的重兒子,重兒子同樹鏈剖分一樣,是擁有節點最多子樹的兒子,這個過程顯然可以O(n)完成

​ 我們用check[i]表示顏色 i 有沒有出現過,ans[i]表示他的顏色個數

​ 遍歷一個節點,我們按以下的步驟進行遍歷:

  • 先遍歷其非重兒子,獲取它的ans,但不保留

  • 遍歷後它的check遍歷它的重兒子,保留它的check

  • 再次遍歷其非重兒子及其父親,用重兒子的check對遍歷到的節點進行計算,獲取整棵子樹的ans

    為什麼不合並第一步和第三步呢?因為check陣列不能重複使用,否則空間會太大,需要在 O(n) 的空間內完成。

    程式碼
    #include<bits/stdc++.h>
    using namespace std;
    #define ll int
    #define r register 
    #define A 1001010
    ll head[A],nxt[A],ver[A],size[A],col[A],cnt[A],ans[A],son[A];
    ll tot=0,num,sum,nowson,n,m,xx,yy;
    inline void add(ll x,ll y){
        nxt[++tot]=head[x],head[x]=tot,ver[tot]=y;
    }
    inline ll read(){
        ll f=1,x=0;char c=getchar();
        while(!isdigit(c)){
            if(c=='-') f=-1;
            c=getchar();
        }
        while(isdigit(c))
            x=(x<<1)+(x<<3)+(c^48),c=getchar();
        return f*x;
    }
    void dfs(ll x,ll fa){
        size[x]=1;
        for(ll i=head[x];i;i=nxt[i]){
            ll y=ver[i];
            if(y==fa) continue;
            dfs(y,x);
            size[x]+=size[y];
            if(size[son[x]]<size[y])
                son[x]=y;
        }
    }
    void cal(ll x,ll fa,ll val){
        if(!cnt[col[x]]) ++sum;
        cnt[col[x]]+=val;
        for(ll i=head[x];i;i=nxt[i]){
            ll y=ver[i];
            if(y==fa||y==nowson) continue;
            cal(y,x,val); 
        }
    }
    void dsu(ll x,ll fa,bool op){
        for(ll i=head[x];i;i=nxt[i]){
            ll y=ver[i];
            if(y==fa||y==son[x])
                continue;
            dsu(y,x,0);
            //從輕兒子出發
        }
        if(son[x])
            dsu(son[x],x,1),nowson=son[x];
        cal(x,fa,1);nowson=0;
        ans[x]=sum;
        if(!op){
            cal(x,fa,-1);
            sum=0;
        }
    }
    int main(){
        n=read();
        for(ll i=1;i<=n-1;i++){
            xx=read(),yy=read();
            add(xx,yy),add(yy,xx);
        }
        for(ll i=1;i<=n;i++)
            col[i]=read();
        dfs(1,0);
        dsu(1,0,1);
        m=read();
        for(ll i=1;i<=m;i++){
            xx=read();
            printf("%d\n",ans[xx]);
        }
    }
    

例題2

CF600E Lomsat gelral
思路

這道題我們可以遍歷整棵樹,並用一個數組ap(appear)記錄每種顏色出現幾次

先用一遍dfs算出每個點是否為重兒子

再dfs統計答案,每次碰到重兒子就跳過,遞迴完清空ap陣列等東東

最後dfs重兒子,不清空

再對當前節點進行另一種dfs,暴力統計ap,不做重兒子

程式碼
#include<bits/stdc++.h>
#define ll long long
#define re register
using namespace std;
const int N=2e5+10;
int n;
int c[N];//color
int v[N],nex[N],first[N],tot=1;
inline void add(int x,int y){
	v[++tot]=y;
	nex[tot]=first[x];
	first[x]=tot;
}
inline int read(){
	int x=0;char ch=getchar();
	while(!isdigit(ch))ch=getchar();
	while(isdigit(ch))x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
	return x;
}
ll ans[N],ap[N],mx,sum;//十年OI一場空,不開 long long 見祖宗 
//ap表示每種顏色出現幾次 mx表示出現最多的次數 sum表示顏色編號和 
int sz[N];//子樹大小 
bool gson[N];//表示一個點是否為重兒子 
void getg(int x,int f){//get 子樹大小 以及 重兒子 
	sz[x]=1;
	int mx=0,p=0;
	for(re int i=first[x];i;i=nex[i]){
		int y=v[i];
		if(y==f)continue;
		getg(y,x);sz[x]+=sz[y];
		if(sz[y]>mx){
			mx=sz[y];
			p=y;
		}
	}if(p)gson[p]=1;
}
void DFS(int x,int f,int p){//暴力遍歷子樹 p為重兒子 之後需init清空 
	//統計答案 
	ap[c[x]]++;
	if(ap[c[x]]>mx){
		mx=ap[c[x]];
		sum=c[x];
	}else if(ap[c[x]]==mx)sum+=c[x];
	for(re int i=first[x];i;i=nex[i]){
		int y=v[i];
		if(y==f || y==p)continue;//不要把重兒子也一起遍歷了! 
		DFS(y,x,p);
	}
}
inline void init(int x,int f){//暴力遍歷後清空 
	ap[c[x]]--;
	for(re int i=first[x];i;i=nex[i]){
		int y=v[i];
		if(y==f)continue;
		init(y,x);
	}
}
void dfs(int x,int f){//啟發式合併關鍵函式! 
	int p=0;//重兒子標記 
	for(re int i=first[x];i;i=nex[i]){
		int y=v[i];
		if(y==f)continue;
		if(!gson[y]){//不是重兒子的暴力做 
			dfs(y,x);
			init(y,x);
			sum=mx=0;
		}
		else p=y;
	}if(p)dfs(p,x);//重兒子單獨特判
	DFS(x,f,p);
	ans[x]=sum;
}
int main(){
	n=read();
	for(re int i=1;i<=n;i++)c[i]=read();
	for(re int i=1;i<n;i++){
		int x=read(),y=read();
		add(x,y),add(y,x);
	}getg(1,0);
	dfs(1,0);
	for(re int i=1;i<=n;i++)
		printf("%lld ",ans[i]);
}

(自己寫的程式WA了,現在只能拿一個題解的程式碼)