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

发布时间 2023-05-05 17:46:04作者: Hamine

解决树上询问问题,没有修改
时间复杂度:\(O(nlogn)\)

例题:https://codeforc.es/contest/600/problem/E
题意:给定一颗树,每个节点有一个颜色,求出子树中颜色最多的颜色值之和。

代码:

#include<bits/stdc++.h>
using namespace std;
using LL = long long;
vector<int> color;
struct HLD{
	vector<int> siz, son, cnt;
	vector<LL> ans;
	vector<vector<int>> e;
	LL sum, Max;
	int hson;
	HLD(int n){
		init(n);
	}
	void init(int n){
		siz.resize(n + 1);
		son.resize(n + 1);
		ans.resize(n + 1);
		cnt.resize(n + 1);
		e.resize(n + 1);
		hson = 0;
		sum = 0;
		Max = 0;
	}
	void add(int u, int v){
		e[u].push_back(v);
		e[v].push_back(u);
	}
	void dfs1(int u, int fa){
		siz[u] = 1;
		for (auto v : e[u]){
			if (v == fa) continue;
			dfs1(v, u);
			siz[u] += siz[v];
			if (siz[v] > siz[son[u]]) son[u] = v;
		}
	}
	void calc(int u, int fa, int val){
		cnt[color[u]] += val;
		if (cnt[color[u]] > Max){
			Max = cnt[color[u]];
			sum = color[u];
		}
		else if (cnt[color[u]] == Max){
			sum += color[u];
		}
		for (auto v : e[u]){
			if (v == fa || v == hson) continue;
			calc(v, u, val);
		}
	}
	void dfs2(int u, int fa, int op){
		for (auto v : e[u]){
			if (v == fa || v == son[u]) continue;
			dfs2(v, u, 0);
		}
		if (son[u]){
			dfs2(son[u], u, 1);
			hson = son[u];	//记录重链编号,计算的时候跳过
		}
		calc(u, fa, 1);
		hson = 0;	//消除的时候所有儿子都清除
		ans[u] = sum;
		if (!op){
			calc(u, fa, -1);
			sum = 0;
			Max = 0;
		}
	}
};
int main(){
	ios::sync_with_stdio(false);cin.tie(0);
	int n;
	cin >> n;
	color.resize(n + 1);
	for (int i = 1; i <= n; i ++ ){
		cin >> color[i];
	}
	HLD t(n);
	for (int i = 0; i < n - 1; i ++ ){
		int u, v;
		cin >> u >> v;
		t.add(u, v);
	}
	t.dfs1(1, 0);
	t.dfs2(1, 0, 0);
	for (int i = 1; i <= n; i ++ ){
		cout << t.ans[i] << " \n"[i == n];
	}
	return 0;
}