WA#3 求助
查看原帖
WA#3 求助
362750
TernaryTree楼主2023/9/25 18:22
#include <bits/stdc++.h>
#define int long long
#define fs first
#define sc second

using namespace std;

const int maxn = 1e5 + 10;
typedef unordered_map<int, int> umap;
typedef pair<int, int> pii;

int n;
int a[maxn];
int siz[maxn];
int son[maxn];
int res[maxn];
int mxd[maxn];
vector<int> g[maxn];

void dfs1(int u, int fa) {
	siz[u] = 1;
	for (int v : g[u]) {
		if (v == fa) continue;
		dfs1(v, u);
		siz[u] += siz[v];
		if (siz[v] > siz[son[u]]) son[u] = v;
	}
}

umap dfs2(int u, int fa) {
	umap ans = umap();
	if (!son[u]) {
		ans[a[u]] = 1;
		mxd[a[u]] = 1;
		res[u] = a[u];
		return ans;
	}
	ans = dfs2(son[u], u);
	int mx = mxd[son[u]], tot = res[son[u]];
	for (int v : g[u]) {
		if (v == fa || v == son[u]) continue;
		umap vm = dfs2(v, u);
		for (pii x : vm) {
			if (ans.count(x.fs)) ans[x.fs] += x.sc;
			else ans[x.fs] = x.sc;
			if (ans[x.fs] > mx) {
				mx = ans[x.fs];
				tot = x.fs;
			} else if (ans[x.fs] == mx) {
				tot += x.fs;
			}
		}
		umap().swap(vm);
	}
	if (ans.count(a[u])) ans[a[u]]++;
	else ans[a[u]] = 1;
	if (ans[a[u]] > mx) {
		mx = ans[a[u]];
		tot = a[u];
	} else if (ans[a[u]] == mx) {
		tot += a[u];
	}
	res[u] = tot;
	mxd[u] = mx;
	return ans;
}

signed main() {
	cin >> n;
	for (int i = 1; i <= n; i++) cin >> a[i];
	for (int i = 1, u, v; i < n; i++) {
		cin >> u >> v;
		g[u].push_back(v);
		g[v].push_back(u);
	}
	dfs1(1, 0);
	dfs2(1, 0);
	for (int i = 1; i <= n; i++) cout << res[i] << " ";
	return 0;
}
2023/9/25 18:22
加载中...