樹上啟發式合併

阿新 • • 發佈：2020-08-16

模板題: CF600E Lomsat gelral

題目描述

一棵樹有n個結點，每個結點都是一種顏色，每個顏色有一個編號，求樹中每個子樹的最多的顏色編號的和。

一些定義

啟發式演算法:啟發式演算法是指基於經驗和直觀感覺，從而對一些演算法的優化。

舉例:並查集的按秩合併

在並查集的按秩合併中,我們將小的集合往大的集合上合併,這樣明顯有利於加快並查集的祖先查詢

演算法流程

首先是一次$bfs$,求出每個節點的重兒子

void dfs1(int u,int fa){
	size[u] = 1 ;//子樹大小
	for(int i = head[u];i;i = edge[i].next){
		int v = edge[i].to;
		if(v==fa) continue;
		dfs1(v,u);
		size[u]+=size[v];
		if(size[v]>size[son[u]]) son[u] = v;//找重兒子
	}
	
}

接下來定義兩個陣列$cnt[]$和$c[]$,分別代表存放的某顏色在“當前”子樹中的數量和存放某節點的顏色

這裡的"當前"指的就是目前正在處理的節點(如果給每個節點都開一個$cnt$的話則會$MLE$)

如果目前正在處理的節點是輕兒子,就把它的答案計入並刪除其貢獻
反之,如果是重兒子,也把它的答案計入,但不刪除其貢獻

void cunt(int u,int fa,int val){
	c[color[u]]+=val;//val為1代表計入貢獻,為-1代表刪除貢獻
	if(c[color[u]] > maxn){//最多的顏色
		maxn = c[color[u]];
		sum = color[u];
	}
	else if(c[color[u]] == maxn){
		sum+=color[u];
	}
	for(int i=head[u];i;i=edge[i].next){
		int v = edge[i].to;
		if(v==maxson||v==fa) continue;//如果是u的重兒子,直接跳過
		cunt(v,u,val);//dfs暴力計貢獻
	}
}
void dfs2(int u,int fa,int keep){//keep代表是否保留該貢獻
	for(int i=head[u];i;i=edge[i].next){
		int v = edge[i].to;
		if(v==son[u]||v==fa) continue;//是重兒子直接跳過
		dfs2(v,u,0);
	}
    
	if(son[u]){//如果有重兒子
		
		dfs2(son[u],u,1);//keep為1,保留其貢獻
		maxson = son[u];//記u節點的重兒子
	}
	cunt(u,fa,1);//暴力統計其非子樹貢獻
	maxson = 0;
	ans[u] = sum;//記錄答案
	if(!keep){//如果不是重兒子,則將其貢獻刪除
		cunt(u,fa,-1); 
		sum = maxn = 0 ;
	}
}

整個$dfs$大致可以分為下面四個流程:

記錄輕兒子及其子樹的答案且刪除其貢獻
記錄重兒子及其子樹的答案且不刪除其貢獻
暴力統計$u$及其所有輕兒子的貢獻合併到剛算出的重兒子資訊裡
刪除該刪除的貢獻

這樣一輪下來相當於是遍歷了兩遍輕兒子,一遍重兒子,顯然效率是較高的

時間複雜度為$O(nlogn)$,~~具體怎麼證還不太清楚~~

$code:$

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int MAXN = 100010;
struct e{
	int next,to;
}edge[MAXN<<1];
int size[MAXN],son[MAXN];
int color[MAXN],c[MAXN];
int maxn , sum;
int head[MAXN<<1],n,cnt = 0;
int ans[MAXN] , maxson;
void add(int u,int v){
      cnt++;
      edge[cnt].to = v;
      edge[cnt].next=head[u];
      head[u]=cnt;
      return;
}
void dfs1(int u,int fa){
	size[u] = 1 ;
	for(int i = head[u];i;i = edge[i].next){
		int v = edge[i].to;
		if(v==fa) continue;
		dfs1(v,u);
		size[u]+=size[v];
		if(size[v]>size[son[u]]) son[u] = v;
	}
	
}
void cunt(int u,int fa,int val){
	c[color[u]]+=val;
	if(c[color[u]] > maxn){
		maxn = c[color[u]];
		sum = color[u];
	}
	else if(c[color[u]] == maxn){
		sum+=color[u];
	}
	for(int i=head[u];i;i=edge[i].next){
		int v = edge[i].to;
		if(v==maxson||v==fa) continue;
		cunt(v,u,val);
	}
}
void dfs2(int u,int fa,int keep){
	for(int i=head[u];i;i=edge[i].next){
		int v = edge[i].to;
		if(v==son[u]||v==fa) continue;
		dfs2(v,u,0);
	}
	if(son[u]){
		
		dfs2(son[u],u,1);
		maxson = son[u];
	}
	cunt(u,fa,1);
	maxson = 0;
	ans[u] = sum;
	if(!keep){
		cunt(u,fa,-1); 
		sum = maxn = 0 ;
	}
}
signed main(){
	scanf("%lld",&n);
	for(int i=1;i<=n;i++){
		scanf("%lld",&color[i]);
	}
	for(int i=1;i<n;i++){
		int u,v;
		scanf("%lld%lld",&u,&v);
		add(u,v);
		add(v,u);
	}
	dfs1(1,0);
	dfs2(1,0,0);
	for(int i=1;i<=n;i++){
		cout<<ans[i]<<" ";
	}
	return 0;
}

~~最近沉迷stg無法自拔了~~