CF 600E 树上启发式合并

#include<cstdio>
#include<vector>
#include<algorithm>
#include<iostream>
#include<set>
using namespace std;
typedef long long ll;
const int maxn = 1e5 + 7;
ll list[maxn];
ll ans, mx;
int siz[maxn];//根
int son[maxn];//重儿子
ll cns[maxn];

vector<int>G[maxn];

void insert(int be, int en) {
	G[be].push_back(en);
}
int  dfs(int x, int fa) {
	siz[x] = 1;
	int c = 0;
	for(int p : G[x]) {
		if (p == fa) continue;
		dfs(p, x);
		siz[x] += siz[p];
		if (c < siz[p]) {
			c = siz[p];
			son[x] = p;
		}
	}
	return 0;
}
ll cnt[maxn];//出现次数

int cal(int x, int fa, int cn, int ban) {
	cnt[list[x]] += cn;

	if (cnt[list[x]] > mx) {
		mx = cnt[list[x]];//times
		ans = list[x];
	}
	else if (cnt[list[x]] == mx) {
		ans += list[x];
	}

	for (int p : G[x]) {
		if (p == fa || p == ban) continue;
		cal(p, x, cn, ban);
	}
	return 0;
}
int dfs2(int x, int fa,int f) {
	for (int p : G[x]) {
		if (p == fa || p == son[x]) continue;
		
		dfs2(p, x, 1);
	}
	if (son[x] ) {
		dfs2(son[x], x, 0);
	}
	cal(x, fa, 1, son[x]);

	cns[x] = ans;

	if (f == 1) {
		cal(x, fa, -1, 0);//把树清理干净
		ans = 0;
		mx = 0;
	}
	return 0;
}


int main() {
	int n;
	int be, en;
	ans = 0;
	mx = 0;
	scanf("%d", &n);
	for (int i = 1; i <= n; i++) {
		scanf("%lld", &list[i]);
	}
	for (int i = 1; i < n; i++) {
		scanf("%d %d", &be, &en);
		insert(be, en);
		insert(en, be);
	}
	dfs(1, -1);
	dfs2(1, -1, 0);
	for (int i = 1; i <= n; i++) {
		if (i == 1) printf("%lld", cns[i]);
		else printf(" %lld", cns[i]);
	}
	printf("\n");
	return 0;
}