codeforces600E Lomsat gelral【線段樹合併/DSU】

  • 2022 年 8 月 13 日
  • 筆記

第一次AC這道題,是三年前的一個下午,也許晚上也說不定。當時使用的\(DSU\) \(on\) \(tree\)算法,如今已經淡忘,再學習新的算法過程中,卻與舊物重逢。生活中充滿不可知會的相遇,即使重逢時多是物是人非。我又想起十六歲時在上海那個船上的夜晚,兩岸的霓虹在無邊黑暗中照亮了她的臉龐,綠衣在少年的眼瞳里凝滯了時間,那真是個美麗的夜晚,即使看不見星月。前天她給我說說點了個贊,我都高興了好久,真是卑微呵。那麼人世中所有明知卑微卻不求告解的人,就是可悲嗎?大抵是有更深沉的緣由的,比如最初的悸動從不屬於初戀的對象,而是屬於動情者的青蔥。詩人和舔狗多少是有些區別的,雖然不多,而對於沿葉節點向根的路途,上面的答案會有下面的影響,所以至下而上一路合併即可。合併的過程中,若是同時抵達葉節點,兩顆原本子樹代表的顏色應該別無二致。樹是權值的,一種顏色對應一個節點,所以\(dfs\)過程中針對顏色單點對應節點修改。至於關鍵的\(push\) \(up\),多者則獨上,相同則齊加,生活亦然。

# include "bits/stdc++.h"
using namespace std;
constexpr int N = 1e5 + 3;
int main() {
	ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
	int n;
	cin >> n;
	vector<int> col(n + 1);
	for(int i = 1; i <= n; ++i) cin >> col[i];
	vector G(n + 1, vector<int>());
	for(int i = 1; i < n; ++i) {
		int u, v;
		cin >> u >> v;
		G[u].push_back(v);
		G[v].push_back(u);
	}
	struct node {
		int l, r, tot;
		long long sum;
	};
	# define ls t[rt].l
	# define rs t[rt].r
	# define lson t[rt].l, l, mid
	# define rson t[rt].r, mid + 1, r
	vector<node> t(N * 50);
	auto push_up = [&](int rt) -> void {
		if(t[ls].tot == t[rs].tot) {
			t[rt].tot = t[ls].tot;
			t[rt].sum = t[ls].sum + t[rs].sum;
		}
		else if(t[ls].tot > t[rs].tot) {
			t[rt].tot = t[ls].tot;
			t[rt].sum = t[ls].sum;
		}
		else {
			t[rt].tot = t[rs].tot;
			t[rt].sum = t[rs].sum;
		}
	};
	vector<int> rt(n + 1, 0);
	int tree_index = n; // 1~n node has their names
	auto update = [&](auto update, int &rt, int l, int r, int x) -> void {
		if(!rt) rt = ++tree_index;
		if(l == r) {
			++t[rt].tot;
			t[rt].sum = l;
			return;
		}
		int mid = l + r >> 1;
		if(x <= mid)
			update(update, lson, x);
		else
			update(update, rson, x);
		push_up(rt);
	};
	auto merge = [&](auto merge, int x, int y, int l, int r) -> int {
		if(!x || !y) return x | y;
		if(l == r) {
			t[x].tot += t[y].tot;
			t[x].sum = l;
			return x;
		}
		int mid = l + r >> 1;
		t[x].l = merge(merge, t[x].l, t[y].l, l, mid);
		t[x].r = merge(merge, t[x].r, t[y].r, mid + 1, r);
		push_up(x);
		return x;
	};
	vector<long long> ans(n + 1);
	auto dfs = [&](auto dfs, int u, int father) -> void {
		for(auto v : G[u]) {
			if(v == father) continue;
			dfs(dfs, v, u);
			merge(merge, u, v, 1, n);
		}
		update(update, u, 1, n, col[u]);
		ans[u] = t[u].sum;
	};
	dfs(dfs, 1, 0);
	for(int i = 1; i <= n; ++i) cout << ans[i] << " ";
	return 0;
}