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
| #include <bits/stdc++.h>
const int IN_LEN = 3000000, OUT_LEN = 400000;
char ibuf[IN_LEN], obuf[OUT_LEN], *ih = ibuf, *oh = obuf;
inline void read(int &x) { for (x = 0; !isdigit(*ih); ih++); while (isdigit(*ih)) x = (x << 1) + (x << 3) + ((*ih++) ^ '0'); }
template<class T>
inline void write(T x) {
static int buf[30], cnt;
if (!x) *oh++ = 48;
else {
for (cnt = 0; x; x /= 10) buf[++cnt] = x % 10 + 48;
while (cnt) *oh++ = buf[cnt--];
}
}
const int MAXN = 300005;
struct Edge {
int to;
Edge *next;
inline void init(const int to, Edge *next) { this->to = to, this->next = next; }
} edge[MAXN << 1], *head[MAXN], *cur = edge;
inline void addEdge(int u, int v) { (++cur)->init(v, head[u]), head[u] = cur; }
int n, croot, csize, S, tmp;
int a[MAXN], size[MAXN];
long long ans[MAXN], tot;
int s[MAXN];
bool vis[MAXN], used[MAXN];
inline void getRoot(int root, int fa) {
register int tmp = 0;
size[root] = 1;
for (Edge *i = head[root]; i; i = i->next)
if ((!vis[i->to]) && (i->to != fa))
getRoot(i->to, root), size[root] += size[i->to], tmp = std::max(tmp, size[i->to]);
if (S - size[root] > tmp) tmp = S - size[root];
if (tmp < csize) croot = root, csize = tmp;
}
inline void dfs1(int root, int fa, int v) {
bool valid = false;
if ((!used[a[root]]) && (a[root] != tmp)) used[a[root]] = valid = true, s[a[root]] += size[root] * v, tot += size[root] * v;
for (Edge *i = head[root]; i; i = i->next) if ((!vis[i->to]) && (i->to != fa)) dfs1(i->to, root, v);
if (valid) used[a[root]] = false;
}
inline void dfs2(int root, int fa, int p) {
bool valid = false;
if ((!used[a[root]]) && (a[root] != tmp)) used[a[root]] = valid = true, tot += p - s[a[root]];
ans[root] += tot;
for (Edge *i = head[root]; i; i = i->next)if ((!vis[i->to]) && (i->to != fa)) dfs2(i->to, root, p);
if (valid) used[a[root]] = false, tot -= p - s[a[root]];
}
inline void solve(int root) {
csize = n, getRoot(root, 0), getRoot(croot, 0), tmp = a[croot];
for (Edge *i = head[croot]; i; i = i->next)if (!vis[i->to]) dfs1(i->to, croot, 1);
tot += size[croot]; ans[croot] += tot;
for (Edge *i = head[croot]; i; i = i->next)
if (!vis[i->to])
dfs1(i->to, croot, -1), tot -= size[i->to], dfs2(i->to, croot, size[croot] - size[i->to]), dfs1(i->to, croot, 1), tot += size[i->to];
for (Edge *i = head[croot]; i; i = i->next) if (!vis[i->to]) dfs1(i->to, croot, -1);
tot = 0, vis[croot] = true;
for (Edge *i = head[croot]; i; i = i->next) if (!vis[i->to]) S = size[i->to], solve(i->to);
}
int main() {
fread(ibuf, 1, IN_LEN, stdin), read(n);
for (register int i = 1; i <= n; i++) read(a[i]);
for (register int i = 1, u, v; i < n; i++) read(u), read(v), addEdge(u, v), addEdge(v, u);
S = n, solve(1);
for (register int i = 1; i <= n; i++) write(ans[i]), *oh++ = '\n';
fwrite(obuf, 1, oh - obuf, stdout);
return 0;
}
|
