Submission #12159656


Source Code Expand

Copy
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

Submission Time
Task F - path pass i
User maspy
Language Python (3.8.2)
Score 600
Code Size 1427 Byte
Status AC
Exec Time 1750 ms
Memory 164128 KB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 600 / 600
Status
AC × 5
AC × 24
Set Name Test Cases
Sample sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt
All line_01.txt, line_3_01.txt, random_01.txt, random_02.txt, random_03.txt, random_100_01.txt, random_100_02.txt, random_10_01.txt, random_10_02.txt, random_1_01.txt, random_1_02.txt, random_2_01.txt, random_2_02.txt, random_2_03.txt, random_3_01.txt, random_3_02.txt, random_3_03.txt, sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt, star_01.txt, star_3_01.txt
Case Name Status Exec Time Memory
line_01.txt AC 1052 ms 136468 KB
line_3_01.txt AC 843 ms 128216 KB
random_01.txt AC 1736 ms 164128 KB
random_02.txt AC 1750 ms 163432 KB
random_03.txt AC 1741 ms 160800 KB
random_100_01.txt AC 1449 ms 145816 KB
random_100_02.txt AC 1453 ms 146132 KB
random_10_01.txt AC 1394 ms 136196 KB
random_10_02.txt AC 1428 ms 136132 KB
random_1_01.txt AC 1364 ms 134212 KB
random_1_02.txt AC 1329 ms 133516 KB
random_2_01.txt AC 1396 ms 133256 KB
random_2_02.txt AC 1365 ms 133504 KB
random_2_03.txt AC 1353 ms 133644 KB
random_3_01.txt AC 1371 ms 134136 KB
random_3_02.txt AC 1389 ms 134080 KB
random_3_03.txt AC 1350 ms 133768 KB
sample_01.txt AC 19 ms 9416 KB
sample_02.txt AC 24 ms 9188 KB
sample_03.txt AC 19 ms 9548 KB
sample_04.txt AC 21 ms 9416 KB
sample_05.txt AC 19 ms 9552 KB
star_01.txt AC 1016 ms 132848 KB
star_3_01.txt AC 887 ms 126040 KB