Problem p02710 (python)

Fixed submission:
import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(505050)
N = int(input())
C = [int(c)-1 for c in input().split()]
G = [[] for _ in range(N)]
for _ in range(N-1):
    a, b = map(int, input().split())
    a -= 1
    b -= 1
    G[a].append(b)
    G[b].append(a)

Cnt = [0] * N
Ans = [0] * N
Siz = [0] * N

def dfs1(v):
    siz = 1
    for u in G[v]:
        G[u].remove(v)
        siz += dfs1(u)
    Siz[v] = siz
    return siz

#Cnt_down = [0] * N

def dfs(v):
    color = C[v]
    cnt_old = Cnt[color]
    siz = 1
    siz_up = N
    #cnt_down = 0
    for u in G[v]:
        #Cnt_down[color] = 0
        Cnt[color] = N - Siz[u]  # 上を塞ぐ
        s = dfs(u)
        siz += s
        siz_up -= s
        Ans[color] += s * (siz_up-cnt_old)
        #cnt_down += Cnt_down[color]
    Ans[color] += siz_up - cnt_old
    #print(siz_up, cnt_old, cnt_down)
    Cnt[color] = cnt_old + siz
    #print(v, Ans)
    return siz

dfs1(0)
dfs(0)
print("\n".join(map(str, Ans)))

Buggy submission (verdict: Runtime Error):
# リストを使わないように書き換え

def numba_compile(numba_config):
    import os, sys
    if sys.argv[-1] == "ONLINE_JUDGE":
        from numba import njit
        from numba.pycc import CC
        cc = CC("my_module")
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature)(func)
            cc.export(func.__name__, signature)(func)
        cc.compile()
        exit()
    elif os.name == "posix":
        exec(f"from my_module import {','.join(func.__name__ for func, _ in numba_config)}")
        for func, _ in numba_config:
            globals()[func.__name__] = vars()[func.__name__]
    else:
        from numba import njit
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature, cache=True)(func)
        print("compiled!", file=sys.stderr)

import numpy as np

def sort_edges(edges_):
    M = len(edges_)
    edges = np.empty((M * 2, 2), dtype=np.int64)
    edges[:M] = edges_
    edges[M:] = edges_[:, ::-1]
    order = np.argsort(edges[:, 0])  # O(N) にできなくもない
    edges = edges[order, 1]
    c = np.cumsum(np.bincount(edges_.ravel()))
    lefts = np.zeros(len(c) + 1, dtype=np.int64)
    lefts[1:] = c
    return edges, lefts

def eular_tour(edges, lefts, root):
    # グラフは 1-indexed が良い
    n = len(lefts)-1
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  # 帰りがけの辺の表現をマイナス以外にしたい場合ここを変える
            continue
        p = parents[v]
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in edges[lefts[v]:lefts[v+1]]:
            if p != u:
                parents[u] = v
                stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

def solve(N, C, AB):
    edges, lefts = sort_edges(AB)
    tour, firsts, lasts, parents = eular_tour(edges, lefts, 1)
    Siz = lasts - firsts + 1 >> 1
    Cnt = np.zeros(N+1, dtype=np.int64)
    Ans = np.zeros(N+1, dtype=np.int64)
    Cnt_old = np.empty_like(Cnt)
    Siz_tmp = np.full_like(Cnt, 1)
    for v in tour:
        if v >= 0:
            p = parents[v]
            if p != -100:
                Cnt[C[p]] = N - Siz[v]  # 上を塞ぐ

            color = C[v]
            Cnt_old[v] = Cnt[color]
        else:
            v = -v
            Ans[C[v]] += N - Siz_tmp[v] + 1 - Cnt_old[v]
            Cnt[C[v]] = Cnt_old[v] + Siz_tmp[v]

            p = parents[v]
            if p != -100:
                Siz_tmp[p] += Siz[v]
                Ans[C[p]] += Siz[v] * (N - Siz_tmp[p] + 1 - Cnt_old[p])

    return Ans[1:]


numba_compile([
    [sort_edges, "Tuple((i8[:],i8[:]))(i8[:,:])"],
    [eular_tour, "Tuple((i8[:],i8[:],i8[:],i8[:]))(i8[:],i8[:],i8)"],
    [solve, "i8[:](i8,i8[:],i8[:,:])"],
])


import sys
def main():
    input = sys.stdin.buffer.readline
    read = sys.stdin.buffer.read
    N = int(input())
    C = np.empty(N+1, dtype=np.int64)
    C[1:] = input().split()
    AB = np.array(read().split(), dtype=np.int64).reshape(N-1, 2)
    Ans = solve(N, C, AB)
    print("\n".join(map(str, Ans.tolist())))

main()