Problem p02710 (python)

Fixed submission:
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])

Buggy submission (verdict: Runtime Error):
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[:])')
    cc.compile()


if sys.argv[-1] == 'ONLINE_JUDGE':
    cc_export()
    exit()

from my_module import solve
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)))