Submission #12198896


Source Code Expand

Copy
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np


def solve(N, C, AB):
    def edges_to_array(A, B, N, M):
        head = np.full(N + 1, -1, np.int32)
        nxt = np.empty(M + M, np.int32)
        to = np.empty(M + M, np.int32)
        for i in range(M):
            a = A[i]
            b = B[i]
            nxt[i << 1] = head[a]
            to[i << 1] = b
            head[a] = i << 1
            nxt[i << 1 | 1] = head[b]
            to[i << 1 | 1] = a
            head[b] = i << 1 | 1
        return head, nxt, to

    def EulerTour(head, nxt, to, root=1):
        N = len(head) - 1
        parent = np.zeros_like(head)
        ind_L = np.empty_like(head)
        ind_R = np.empty_like(head)
        tour = np.empty(N + N, np.int32)
        stack = np.empty(N + N, np.int32)
        stack[0] = -root
        stack[1] = root
        t = -1
        s = 2
        while s > 0:
            s -= 1
            v = stack[s]
            t += 1
            tour[t] = v
            if v > 0:
                ind_L[v] = t
                p = parent[v]
                k = head[v]
                while k != -1:
                    w = to[k]
                    if w == p:
                        k = nxt[k]
                        continue
                    parent[w] = v
                    stack[s] = -w
                    s += 1
                    stack[s] = w
                    s += 1
                    k = nxt[k]
            else:
                ind_R[-v] = t
        return tour, ind_L, ind_R, parent

    head, nxt, to = edges_to_array(AB[::2], AB[1::2], N, N - 1)
    tour, ind_L, ind_R, parent = EulerTour(head, nxt, to)
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    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
    return answer


def cc_export():
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('solve', '(i8,i4[:],i4[:])')(solve)
    cc.compile()


try:
  from my_module import solve
except:
  cc_export()
  exit()
N = int(readline())
C = np.zeros(N + 1, np.int32)
C[1:] = np.array(readline().split(), np.int32)
AB = np.array(read().split(), np.int32)
answer = solve(N, C, AB)
print('\n'.join(answer[1:].astype(str)))

Submission Info

Submission Time
Task F - path pass i
User maspy
Language Python (3.8.2)
Score 600
Code Size 2819 Byte
Status
Exec Time 398 ms
Memory 76384 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 376 ms 76104 KB
line_3_01.txt 356 ms 76064 KB
random_01.txt 394 ms 76320 KB
random_02.txt 398 ms 76384 KB
random_03.txt 397 ms 76252 KB
random_100_01.txt 381 ms 76100 KB
random_100_02.txt 384 ms 76072 KB
random_10_01.txt 373 ms 76208 KB
random_10_02.txt 381 ms 75988 KB
random_1_01.txt 379 ms 75988 KB
random_1_02.txt 372 ms 76204 KB
random_2_01.txt 378 ms 76112 KB
random_2_02.txt 385 ms 75856 KB
random_2_03.txt 370 ms 76212 KB
random_3_01.txt 374 ms 76000 KB
random_3_02.txt 373 ms 75992 KB
random_3_03.txt 379 ms 75844 KB
sample_01.txt 102 ms 27380 KB
sample_02.txt 104 ms 27140 KB
sample_03.txt 102 ms 27020 KB
sample_04.txt 100 ms 27120 KB
sample_05.txt 108 ms 27112 KB
star_01.txt 365 ms 69808 KB
star_3_01.txt 344 ms 70056 KB