Submission #12151858

---

Source Code Expand

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

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(v, p):
    A = dp[v]
    B = dp[p]
    size_A = size[v]
    size_B = size[p]
    c = C[p]
    if len(A) < len(B):
        A, B = B, A
        size_A, size_B = size_B, size_A
    answer[c] += (size_A - A[c]) * (size_B - B[c])
    # B を A にマージ
    for key, cnt in B.items():
        answer[key] -= cnt * A[key]
        A[key] += cnt
    size[p] = size_A + size_B
    dp[p] = A


answer = [0] * (N + 1)
dp = [defaultdict(int) for _ in range(N + 1)]
size = [0] * (N + 1)

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

A = dp[root]
c = size[root]
for v in range(1, N + 1):
    print(c * A[v] + answer[v])

Submission Info

Judge Result