TLE求助(dsu on tree)
查看原帖
TLE求助(dsu on tree)
530349
天空即为极限楼主2021/12/25 17:15

rt

#include<bits/stdc++.h>
#define int long long
using namespace std;
vector<int>v[1000005];
int cnt[2000005], Mx, Son, sum, ans[1000005];
struct node{
	int fa, size, son, color;
}k[1000005];
inline int read() {
    char c = getchar(); int x = 0, f = 1;
    while(c < '0' || c > '9') {if(c == '-') f = -1; c = getchar();}
    while(c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
    return x * f;
}
void dfs(int x, int fa){
	k[x].fa = fa;
	k[x].size = 1;
	for(int i = 0; i < v[x].size(); i++){
		if(v[x][i] == fa) continue;
		dfs(v[x][i], x);
		k[x].size += k[v[x][i]].size;
		if(k[v[x][i]].size > k[k[x].son].size) k[x].son = v[x][i];
	}
}

void add(int x, int fa, int val){
	cnt[k[x].color] += val;
	if(cnt[k[x].color] > Mx) Mx = cnt[k[x].color], sum = k[x].color;
	else if(cnt[k[x].color] == Mx) sum += k[x].color;
	for(int i = 0; i < v[x].size(); i++){
		if(v[x][i] == fa or v[x][i] == Son) continue;
		add(v[x][i], x, val);
	}
}

void dfs2(int x, int fa, bool com){
	for(int i = 0; i < v[x].size(); i++){
		if(v[x][i] == fa) continue;
		if(v[x][i] != k[x].son) dfs2(v[x][i], x, false);
	}
	if(k[x].son != 0) dfs2(k[x].son, x, true), Son = k[x].son;
	add(x, fa, 1); Son = 0;
	ans[x] = sum;
	if(!com) add(x, fa, -1), sum = 0, Mx = 0;
}
signed main(){
	//ios::sync_with_stdio(false);
	//cin.tie(0), cout.tie(0);
	int n, x, y;
	n = read();
	for(int i = 1; i <= n; i++) k[i].color = read();
	for(int i = 1; i <= n - 1; i++){
		x = read(), y = read();
		v[x].push_back(y);
		v[y].push_back(x);
	}
	dfs(0, 1);	
	dfs2(1, 0, false);
	for(int i = 1; i <= n; i++) cout << ans[i] << " ";
}
2021/12/25 17:15
加载中...