mxqz dsu on tree
查看原帖
mxqz dsu on tree
300098
cmaths楼主2022/10/5 12:22
// Problem: CF600E Lomsat gelral
// Contest: Luogu
// URL: https://www.luogu.com.cn/problem/CF600E
// Memory Limit: 250 MB
// Time Limit: 2000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

#include <cstdio>
#include <iostream>
#include <algorithm>
#define int long long

using namespace std;

const int N = 100000;
int n;
int col[N + 5], ans[N + 5];
struct Edge
{
	int u, v, nxt;
}edge[N * 2 + 5];
int head[N + 5], cntE;
void addE(int u, int v)
{
	edge[++cntE] = (Edge){u, v, head[u]};
	head[u] = cntE;
}
int hson[N + 5], siz[N + 5];
void dfs(int u, int f)
{
	siz[u] = 1;
	for(int i = head[u]; i; i = edge[i].nxt)
	{
		int v = edge[i].v;
		if(v == f)
		{
			return;
		}
		dfs(v, u);
		siz[u] += siz[v];
		if(siz[v] > siz[hson[u]])
		{
			hson[u] = v;
		}
	}
}
int flag;
int cnt[N + 5];
int maxc, sum;
void calc(int u, int f, int val)
{
	cnt[col[u]] += val;
	if(cnt[col[u]] > maxc)
	{
		maxc = cnt[col[u]];
		sum = col[u];
	}
	else if(cnt[col[u]] == maxc)
	{
		sum += col[u];
	}
	for(int i = head[u]; i; i = edge[i].nxt)
	{
		int v = edge[i].v;
		if(v == f || v == flag)
		{
			continue;
		}
		calc(v, u, val);
	}
}
void dot(int u, int f, bool kep)
{
	for(int i = head[u]; i; i = edge[i].nxt)
	{
		int v = edge[i].v;
		if(v == f || v == hson[u])
		{
			continue;
		}
		dot(v, u, 0);
	}
	if(hson[u])
	{
		dot(hson[u], u, 1);
		flag = hson[u];
	}
	calc(u, f, 1);
	flag = 0;
	ans[u] = sum;
	if(!kep)
	{
		calc(u, f, -1);
		maxc = sum = 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);
		addE(u, v);
		addE(v, u);
	}
	dfs(1, 0);
	dot(1, 0, 0);
	for(int i = 1; i <= n; i++)
	{
		printf("%lld ", ans[i]);
	}
	return 0;
}

它会在 #18 TLE

2022/10/5 12:22
加载中...