Codeforces 600E Lomsat gelral （Dsu On the Tree）

 1 #include <bits/stdc++.h>
 2 
 3 using namespace std;
 4 
 5 #define rep(i, a, b)    for (int i(a); i <= (b); ++i)
 6 
 7 typedef long long LL;
 8 
 9 const int N = 600010;
10 
11 int n;
12 int cc[N], col[N], sz[N], son[N];
13 LL ans[N];
14 
15 vector <int> v[N];
16 
17 void getsize(int x, int fa){
18         sz[x] = 1;
19         for (auto u : v[x]){
20                 if (u == fa) continue;
21                 getsize(u, x); sz[x] += sz[u];
22                 if (!son[x] || sz[u] > sz[son[x]]) son[x] = u;
23         }
24 }
25 
26 bool skip[N];
27 LL sum = 0;
28 int cx = 0;
29 int x, y;
30 
31 void edt(int x, int fa, int k){
32         cc[col[x]] += k;
33         if (k > 0 && cc[col[x]] >= cx){
34                 if (cc[col[x]] > cx)
35                         sum = 0, cx = cc[col[x]];
36                 sum += col[x];
37         }
38         
39         for (auto u : v[x]){
40         if (u == fa || skip[u]) continue;
41         edt(u, x, k);
42     }
43     
44 }
45 
46 void dfs(int x, int fa, int kep = 0){
47         for (auto u : v[x]){
48                 if (u == fa || u == son[x]) continue;
49                 dfs(u, x, 0);
50         }
51         
52         if (son[x]) dfs(son[x], x, 1), skip[son[x]] = 1;
53         edt(x, fa, 1);
54         
55         ans[x] = sum;
56         if (son[x]) skip[son[x]] = 0;
57         if (!kep) edt(x, fa, -1), cx = sum = 0;
58 }
59 
60 int main(){
61     
62     scanf("%d", &n);
63     rep(i, 1, n) scanf("%d", col + i);
64     
65     rep(i, 1, n - 1){
66         scanf("%d%d", &x, &y);
67         v[x].push_back(y);
68         v[y].push_back(x);
69     }
70     
71     getsize(1, 0);
72     dfs(1, 0, 0);
73     
74     rep(i, 1, n) printf("%lld\n", ans[i]);
75     return 0;
76 }