树上启发式合并学习笔记

树上启发式合并 \((dsu\ on \ tree)\)

适用条件：

可以在一个子树内统计的问题，并且不带修改。暴力复杂度一般为 \(O(n^2)\)。

例题：

CF600E Lomsat gelral

解法

考虑一个问题 ，给你一棵树，每个节点有一个颜色，如果一种颜色在以 \(x\) 为根的子树内出现次数最多，则称 \(col_i\) 占主要色。求算对于每一个 \(i∈[1,n]\)，求以 \(i\) 为根的子树中，占主导地位颜色编号和。

我们考虑 dsu on tree。首先维护出每个节点的重儿子 \(son_i\)，而后我们考虑写这样一个 dfs 函数求答案。

首先对于 \(u\) 为根的子树遍历所有的轻儿子，再遍历所有的重儿子，这样就可以求得 \(u\) 子树内所有答案了。这时候我们考虑求算 \(u\) 点的答案，这时候我们求算过的所有的轻儿子答案已经被清空掉了，只剩下子树内重儿子的答案 \(sum\)。这时候我们考虑只暴力遍历轻儿子加入答案即可。

复杂度为 \(O(n\log n)\)。

代码

#include<bits/stdc++.h>
using namespace std;
#define int long long

const int N=5e5+7;
int n,col[N];
vector<int> G[N];
int sz[N],dep[N],son[N];
int num[N],sum,maxN;
int ans[N];

void dfs(int u,int father){
	sz[u]=1;
	for(int k:G[u]){
		if(k==father) continue;
		dfs(k,u);sz[u]+=sz[k];
		if(!son[u]||sz[k]>sz[son[u]]) son[u]=k; 
	}
}

void updata(int u,int father,int son,int val){
	num[col[u]]+=val;
	if(num[col[u]]>maxN) maxN=num[col[u]],sum=col[u];
	else if(num[col[u]]==maxN) sum+=col[u];
	for(int k:G[u]){
		if(k!=father&&k!=son) updata(k,u,son,val);
	}
}

void DFS(int u,int father,bool kep){
	for(int k:G[u]){
		if(k!=father&&k!=son[u]) DFS(k,u,0);
	}
	if(son[u]) DFS(son[u],u,1);
	updata(u,father,son[u],1);
	ans[u]=sum;
	if(!kep) updata(u,father,-1,-1),sum=maxN=0;
}

signed main(){
	scanf("%lld",&n);
	for(int i=1;i<=n;i++) scanf("%lld",&col[i]);
	for(int i=1;i<n;i++){
		int u,v;scanf("%lld%lld",&u,&v);G[u].push_back(v),G[v].push_back(u);
	}
	dfs(1,-1);DFS(1,-1,1);
	for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
	return 0;
}

CF208E Blood Cousins

解法

对于一棵树，或者一座森林，询问 \((v,p)\) 表示编号为 \(v\) 的人有多少个 \(p\) 级表亲。

我们考虑 dsu on tree。首先我们从 \(v\) 跳到 \(p\) 级祖先 \(u\)，然后我们考虑在 \(u\) 所在子树内统计深度为 \(dep[u]+p\) 的点数量和。dsu on tree 思路同上一题。

注意：

• 我们建一个虚拟根 \(0\) 号点便于我们进行 dfs，但是我们在判断是否有解的情况时不能把 \(0\) 当作 \(v\) 的 \(p\) 级祖先，因而我们判断时需要满足 dep[x]>=+1 才有解，反之无解。
• 注意计算贡献或删除贡献时不要重复计算或删除。
• 因为对于每个子树，\(v\) 也被计算过一次，因而最后答案要减一。

复杂度 \(O(n\log n)\)。

代码

#include<bits/stdc++.h>
using namespace std;

const int N=2e6+10;
int n;
vector<int> G[N];
vector<pair<int,int> > Q[N];
int sz[N],dep[N],son[N];
int cnt[N];
int ans[N];
int f[100000][30];
int q;

void dfs(int u,int father){
	f[u][0]=father,dep[u]=dep[father]+1;
	sz[u]=1;
	for(int k:G[u]){
		if(k==father) continue;
		dfs(k,u);sz[u]+=sz[k];
		if(!son[u]||sz[k]>sz[son[u]]) son[u]=k; 
	}
}

void updata(int u,int father,int son,int val){
	cnt[dep[u]]+=val;
	for(int k:G[u]){
		if(k!=father&&k!=son) updata(k,u,son,val);
	}
}

void DFS(int u,int father,bool kep){
	for(int k:G[u]){
		if(k!=father&&k!=son[u]) DFS(k,u,0);
	}
	if(son[u]) DFS(son[u],u,1);
	updata(u,father,son[u],1);
	for(auto it:Q[u]){
		int _dep=it.first,_id=it.second;
		ans[_id]=cnt[_dep];
	}
	if(!kep) updata(u,father,-1,-1);
}

signed main(){
	scanf("%d",&n);
	for(int i=1;i<=n;i++){
		int fa;scanf("%d",&fa);
		G[fa].push_back(i);
	}
	dfs(0,-1);
	for(int j=1;j<=20;j++){
		for(int i=1;i<=n;i++){
			f[i][j]=f[f[i][j-1]][j-1];
		}
	}
	scanf("%d",&q);
	for(int i=1;i<=q;i++){ 
		int x,d;
		scanf("%d%d",&x,&d);
		if(dep[x]<=d+1) continue;
		for(int k=0;k<=20;k++){
			if((1<<k)&d) x=f[x][k];
		}
		Q[x].push_back(make_pair(dep[x]+d,i));//亲戚的深度和问题编号。 
	}
	DFS(0,-1,-1);
	for(int i=1;i<=q;i++) printf("%d ",ans[i]==0?0:ans[i]-1);
	return 0;
}

CF1009F Dominant Indices

解法

对于一个以 \(u\) 为根的子树，定义 \(d(u,k)\) 表示距离 \(u\) 距离为 \(k\) 的节点个数，则求在 \(d(u,k)\to \max\) 时，\(k\) 的最小值。

我们考虑 dsu on tree。维护 \(maxN_i\) 表示深度 \(i\) 的节点个数最大值。维护 \(Dis\) 表示 \(d(u,k)+dep_u\) 的最小值。别忘了最后结果要减掉 \(dep_u\)。

复杂度 \(O(n\log n)\)。

代码

#include<bits/stdc++.h>
using namespace std;
#define int long long

const int N=2e6+7;

int n,col[N];
vector<int> G[N];
int sz[N],dep[N],son[N];
int num[N],sum,maxN;
int ans[N];
int Dis;

void dfs(int u,int father){
	sz[u]=1;dep[u]=dep[father]+1; 
	for(int k:G[u]){
		if(k==father) continue;
		dfs(k,u);sz[u]+=sz[k];
		if(!son[u]||sz[k]>sz[son[u]]) son[u]=k; 
	}
}

void updata(int u,int father,int son,int val){
	num[dep[u]]+=val;
	if(num[dep[u]]>maxN||((maxN==num[dep[u]])&&(dep[u]<Dis))) Dis=dep[u],maxN=num[dep[u]];
	for(int k:G[u]){
		if(k!=father&&k!=son) updata(k,u,son,val);
	}
}

void DFS(int u,int father,bool kep){
	for(int k:G[u]){
		if(k!=father&&k!=son[u]) DFS(k,u,0);
	}
	if(son[u]) DFS(son[u],u,1);
	updata(u,father,son[u],1);
	ans[u]=Dis-dep[u];
	if(!kep) updata(u,father,-1,-1),Dis=0,maxN=0;
}

signed main(){
	scanf("%lld",&n);
	for(int i=1;i<n;i++){
		int u,v;scanf("%lld%lld",&u,&v);G[u].push_back(v),G[v].push_back(u);
	}
	dfs(1,0);DFS(1,0,1);
//	for(int i=1;i<=n;i++) printf("%d ",dep[i]);puts("");
	for(int i=1;i<=n;i++) printf("%lld\n",ans[i]);
	return 0;
}

CF375D Tree and Queries

解法

给定一棵树，每个节点上有一个颜色。\(m\) 次询问，给定一个询问点对 \((u,k)\) 表示在以 \(u\) 为根的子树中，出现次数 \(\ge k\) 的颜色种类数。

我们考虑 dsu on tree。对于每个颜色维护一个出现次数 \(cnt_{col}\) 和 出现次数大于等于 \(k\) 的颜色种类数 \(times_{cnt}\)。每次更新即可。

注意这道题 updata 不能写在一起，因为加的时候我们是该颜色出现次数 +1，更改后的达标次数 +1，而减的时候需要先把达标次数清空或 -1，再把该颜色出现次数 -1。这样避免了清错的情况发生。

复杂度 \(O(n\log n)\)。

代码1

#include<bits/stdc++.h>
using namespace std;
#define int long long

const int N=2e6+7;

int n,col[N];
vector<int> G[N];
int sz[N],dep[N],son[N];
int num[N],sum,maxN;
int ans[N];
int Dis;
int m;
vector<pair<int,int> > Q[N];
int times[N];

void dfs(int u,int father){
	sz[u]=1;dep[u]=dep[father]+1; 
	for(int k:G[u]){
		if(k==father) continue;
		dfs(k,u);sz[u]+=sz[k];
		if(!son[u]||sz[k]>sz[son[u]]) son[u]=k; 
	}
}

void updata(int u,int father,int son,int val){
	num[col[u]]+=val;
	times[num[col[u]]]+=val;
	for(int k:G[u]){
		if(k!=father&&k!=son) updata(k,u,son,val);
	}
}

void clear(int u,int father,int son,int val){
	times[num[col[u]]]=0;
	num[col[u]]--;
	for(int k:G[u]){
		if(k!=father&&k!=son) clear(k,u,son,val);
	}
}

void DFS(int u,int father,bool kep){
	for(int k:G[u]){
		if(k!=father&&k!=son[u]) DFS(k,u,0);
	}
	if(son[u]) DFS(son[u],u,1);
	updata(u,father,son[u],1);
	for(auto it:Q[u]){
		int k=it.first,id=it.second;
		ans[id]=times[k];
	}
	if(!kep) clear(u,father,-1,-1);
}

signed main(){
	scanf("%lld%lld",&n,&m);
	for(int i=1;i<=n;i++) scanf("%lld",&col[i]);
	for(int i=1;i<n;i++){
		int u,v;scanf("%lld%lld",&u,&v);G[u].push_back(v),G[v].push_back(u);
	}
	for(int i=1;i<=m;i++){
		int x,k;
		scanf("%lld%lld",&x,&k);
		Q[x].push_back(make_pair(k,i)); 
	}
	dfs(1,0);DFS(1,0,1);
	for(int i=1;i<=m;i++) printf("%lld\n",ans[i]);
	return 0;
}

代码2

#include<bits/stdc++.h>
using namespace std;
#define int long long

const int N=2e6+7;

int n,col[N];
vector<int> G[N];
int sz[N],dep[N],son[N];
int num[N],sum,maxN;
int ans[N];
int times[N];
int Dis;
int m;
vector<pair<int,int> > Q[N];
int dfn_clock;
int dfn[N];
int pos[N];

void dfs(int u,int father){
	sz[u]=1;dep[u]=dep[father]+1;
	dfn[u]=++dfn_clock;
	pos[dfn_clock]=u;
	for(int k:G[u]){
		if(k==father) continue;
		dfs(k,u);sz[u]+=sz[k];
		if(!son[u]||sz[k]>sz[son[u]]) son[u]=k;
	}
}

void add(int x){
	for(int i=dfn[x];i<=dfn[x]+sz[x]-1;i++){
		num[col[pos[i]]]++;
		times[num[col[pos[i]]]]++;
	}
}

void del(int x){
	for(int i=dfn[x];i<=dfn[x]+sz[x]-1;i++){
		times[num[col[pos[i]]]]=0; 
		num[col[pos[i]]]--;
	}
}

void DFS(int u,int father,bool kep){
	for(int k:G[u]){
		if(k!=father&&k!=son[u]) DFS(k,u,0);
	}
	if(son[u]) DFS(son[u],u,1);
	for(int k:G[u]){
		if(k!=father&&k!=son[u]) add(k);
	}
	
	num[col[u]]++;
	times[num[col[u]]]++;
//	printf("%lld ",times[num[col[u]]]);
	for(auto it:Q[u]){
		int k=it.first,id=it.second;
		ans[id]=times[k];
	}
	if(!kep) del(u);
}

signed main(){
	scanf("%lld%lld",&n,&m);
	for(int i=1;i<=n;i++) scanf("%lld",&col[i]);
	for(int i=1;i<n;i++){
		int u,v;scanf("%lld%lld",&u,&v);G[u].push_back(v),G[v].push_back(u);
	}
	for(int i=1;i<=m;i++){
		int x,k;
		scanf("%lld%lld",&x,&k);
		Q[x].push_back(make_pair(k,i)); 
	}
	dfs(1,0);DFS(1,0,1);
	for(int i=1;i<=m;i++) printf("%lld\n",ans[i]);
	return 0;
}

CF246E Blood Cousins Return

解法

给定一片森林，询问编号为 \(u\) 的 \(k\) 级儿子有多少不同的名字。

考虑 dsu on tree的做法。用 set 维护字符串出现次数。注意读入输出格式。同时注意本题是森林，需要多次 dfs 并且每次 dfs 起始点都不同，都是当前状态下没有 \(fa\) 的点。

复杂度 \(O(n\log n)\)。

代码

#include<bits/stdc++.h>
using namespace std;
#define int long long

const int N=1e6+7;

int n,col[N];
vector<int> G[N];
int sz[N],dep[N],son[N];
int num[N],sum,maxN;
int ans[N];
int fa[N];
int Dis;
int m;
vector<pair<int,int> > Q[N];
int times[N];
string a[N];
set<string> s[N];

void dfs(int u,int father){
	sz[u]=1;dep[u]=dep[father]+1;
	for(int k:G[u]){
		if(k==father) continue;
		dfs(k,u);sz[u]+=sz[k];
		if(!son[u]||sz[k]>sz[son[u]]) son[u]=k; 
	}
}

void updata(int u,int father,int son,int val){
	s[dep[u]].insert(a[u]);
	for(int k:G[u]){
		if(k!=father&&k!=son) updata(k,u,son,val);
	}
}

void clear(int u,int father,int son,int val){
	s[dep[u]].clear();
	for(int k:G[u]){
		if(k!=father&&k!=son) clear(k,u,son,val);
	}
}

void DFS(int u,int father,bool kep){
	for(int k:G[u]){
		if(k!=father&&k!=son[u]) DFS(k,u,0);
	}
	if(son[u]) DFS(son[u],u,1);
	updata(u,father,son[u],1);
	for(auto it:Q[u]){
		int k=it.first,id=it.second;
		ans[id]=s[k].size();
	}
	if(!kep) clear(u,father,-1,-1);
}

signed main(){
	scanf("%lld",&n);
	for(int i=1;i<=n;i++){
		int u;
		cin>>a[i]>>u;
		G[u].push_back(i);
		fa[i]=u;
	}
	for(int i=1;i<=n;i++){
		if(!fa[i]) dfs(i,0);
	}
	scanf("%lld",&m);
		for(int i=1;i<=m;i++){
		int x,k;
		scanf("%lld%lld",&x,&k);
		Q[x].push_back(make_pair(k+dep[x],i)); 
	}
	for(int i=1;i<=n;i++) if(!fa[i]) DFS(i,0,0);
	for(int i=1;i<=m;i++) printf("%lld\n",ans[i]);
	return 0;
}