Submission #12151858


Source Code Expand

Copy
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

Submission Time
Task F - path pass i
User maspy
Language Python (3.8.2)
Score 600
Code Size 1661 Byte
Status
Exec Time 1935 ms
Memory 162800 KB

Judge Result

Set Name Score / Max Score Test Cases
Sample 0 / 0 sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt
All 600 / 600 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 1201 ms 136292 KB
line_3_01.txt 989 ms 128208 KB
random_01.txt 1906 ms 162800 KB
random_02.txt 1935 ms 162180 KB
random_03.txt 1908 ms 159788 KB
random_100_01.txt 1567 ms 146120 KB
random_100_02.txt 1545 ms 146480 KB
random_10_01.txt 1490 ms 136068 KB
random_10_02.txt 1554 ms 136260 KB
random_1_01.txt 1464 ms 134440 KB
random_1_02.txt 1403 ms 133772 KB
random_2_01.txt 1437 ms 136376 KB
random_2_02.txt 1451 ms 133840 KB
random_2_03.txt 1475 ms 133964 KB
random_3_01.txt 1428 ms 134568 KB
random_3_02.txt 1454 ms 134020 KB
random_3_03.txt 1474 ms 134120 KB
sample_01.txt 19 ms 9260 KB
sample_02.txt 19 ms 9300 KB
sample_03.txt 21 ms 9416 KB
sample_04.txt 18 ms 9264 KB
sample_05.txt 24 ms 9244 KB
star_01.txt 1141 ms 130732 KB
star_3_01.txt 1006 ms 125976 KB