CF 600E Lomsat gelral

题目描述

https://codeforces.com/problemset/problem/600/E

Solution

直接上 $dsu$,时间复杂度 $O(n\log n)$

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
#include <iostream>
#include <cstdio>
#include <cstring>
#define maxn 100010
#define ll long long
using namespace std;

int n, col[maxn];;

struct Edge {
int to, next;
} e[maxn * 2]; int c1, head[maxn];
inline void add_edge(int u, int v) {
e[c1].to = v; e[c1].next = head[u]; head[u] = c1++;
}

int sz[maxn], son[maxn], in[maxn], out[maxn], bl[maxn], c2;
void Dfs(int u, int fa) {
int Max = 0; sz[u] = 1; in[u] = ++c2; bl[c2] = u;
for (int i = head[u]; ~i; i = e[i].next) {
int v = e[i].to; if (v == fa) continue;
Dfs(v, u); sz[u] += sz[v]; if (sz[v] > Max) Max = sz[v], son[u] = v;
} out[u] = c2;
}

int cnt1[maxn], cnt2[maxn], mx; ll sum[maxn];
void add(int u) {
for (int i = in[u]; i <= out[u]; ++i) {
int c = col[bl[i]], cnt = ++cnt1[c];
if (cnt > mx) mx = cnt; ++cnt2[cnt]; sum[cnt] += c;
cnt2[cnt - 1]--; sum[cnt - 1] -= c;
}
}

void del(int u) {
for (int i = in[u]; i <= out[u]; ++i) {
int c = col[bl[i]], cnt = --cnt1[c];
if (--cnt2[cnt + 1] == 0 && cnt + 1 == mx) --mx; sum[cnt + 1] -= c;
cnt2[cnt]++; sum[cnt] += c;
}
}

ll ans[maxn];
void dfs(int u, int fa) {
for (int i = head[u]; ~i; i = e[i].next) {
int v = e[i].to; if (v == fa || v == son[u]) continue ;
dfs(v, u); del(v);
}
if (son[u]) dfs(son[u], u);
for (int i = head[u]; ~i; i = e[i].next) {
int v = e[i].to; if (v == fa || v == son[u]) continue;
add(v);
}

int c = col[u], cnt = ++cnt1[c];
if (cnt > mx) mx = cnt; ++cnt2[cnt]; sum[cnt] += c;
cnt2[cnt - 1]--; sum[cnt - 1] -= c;
ans[u] = sum[mx];
}

int main() { memset(head, -1, sizeof head);
cin >> n;
for (int i = 1; i <= n; ++i) scanf("%d", &col[i]);
for (int i = 1; i < n; ++i) {
int x, y; scanf("%d%d", &x, &y);
add_edge(x, y); add_edge(y, x);
} Dfs(1, 0); dfs(1, 0);
for (int i = 1; i <= n; ++i) printf("%I64d ", ans[i]);
return 0;
}