Submission #12191972


Source Code Expand

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


def EulerTour(graph, root=1):
    N = len(graph) - 1
    parent = [0] * (N + 1)
    stack = [-root, root]
    tour = []
    ind_L = [0] * (N + 1)
    ind_R = [0] * (N + 1)
    i = -1
    while stack:
        v = stack.pop()
        tour.append(v)
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
    return tour, ind_L, ind_R, parent


N = int(readline())
graph = [[] for _ in range(N + 1)]
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

tour, ind_L, ind_R, parent = EulerTour(graph)
removed = [0] * (N + 1)
answer = [N * (N + 1) // 2] * (N + 1)
memo = [0] * (N + 1)

for v in tour:
    if v > 0:
        memo[v] = removed[C[parent[v]]]
    else:
        v = -v
        removed[C[v]] += 1
        p = parent[v]
        x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
        x -= removed[C[p]] - memo[v]
        answer[C[p]] -= x * (x + 1) // 2
        removed[C[p]] += x

for i, x in enumerate(removed):
    x = N - x
    answer[i] -= x * (x + 1) // 2
print('\n'.join(map(str, answer[1:])))

Submission Info

Submission Time
Task F - path pass i
User maspy
Language Python (3.8.2)
Score 600
Code Size 1530 Byte
Status AC
Exec Time 1166 ms
Memory 105740 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 825 ms 105740 KB
line_3_01.txt AC 662 ms 93592 KB
random_01.txt AC 1146 ms 103916 KB
random_02.txt AC 1166 ms 102452 KB
random_03.txt AC 1140 ms 102448 KB
random_100_01.txt AC 1069 ms 98072 KB
random_100_02.txt AC 1053 ms 98072 KB
random_10_01.txt AC 1054 ms 97200 KB
random_10_02.txt AC 1045 ms 96980 KB
random_1_01.txt AC 1041 ms 93748 KB
random_1_02.txt AC 1025 ms 93928 KB
random_2_01.txt AC 1054 ms 95956 KB
random_2_02.txt AC 1037 ms 96084 KB
random_2_03.txt AC 1031 ms 95820 KB
random_3_01.txt AC 1047 ms 96240 KB
random_3_02.txt AC 1028 ms 95256 KB
random_3_03.txt AC 1042 ms 95076 KB
sample_01.txt AC 17 ms 9232 KB
sample_02.txt AC 23 ms 9348 KB
sample_03.txt AC 18 ms 9120 KB
sample_04.txt AC 22 ms 9116 KB
sample_05.txt AC 17 ms 9412 KB
star_01.txt AC 741 ms 103556 KB
star_3_01.txt AC 648 ms 93008 KB