Submission #12159656

---

Source Code Expand

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import defaultdict

N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
AB = zip(m, m)

graph = [[] for _ in range(N + 1)]
for a, b in AB:
    graph[a].append(b)
    graph[b].append(a)

root = 1
parent = [0] * (N + 1)
depth = [0] * (N + 1)  # 根を深さ 1 として
order = []
stack = [root]
depth[root] = 1
while stack:
    x = stack.pop()
    order.append(x)
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x
        stack.append(y)
        depth[y] = depth[x] + 1


def merge(A, B):
    if len(A) < len(B):
        A, B = B, A
    for key, cnt in B.items():
        A[key] += cnt
    return A


dp = [defaultdict(int) for _ in range(N + 1)]  # 部分木内で、消去済の頂点数の分布
size = [0] * (N + 1)  # 部分木の大きさ
answer = [N * (N + 1) // 2] * (N + 1)

for v in order[::-1]:
    size[v] += 1
    p = parent[v]
    dp[v][C[v]] = size[v]
    if not p:
        break
    size[p] += size[v]
    c = C[p]
    x = size[v] - dp[v][c]
    answer[c] -= x * (x + 1) // 2
    dp[p] = merge(dp[p], dp[v])

A = dp[root]
for c in range(1, N + 1):
    x = N - A[c]
    answer[c] -= x * (x + 1) // 2

print(*answer[1:], sep='\n')

Submission Info

Judge Result