Problem p02728 (python)

Fixed submission:
import sys
sys.setrecursionlimit(505050)

class Combination:
    def __init__(self, n_max, mod=10**9+7):
        # O(n_max + log(mod))
        self.mod = mod
        f = 1
        self.fac = fac = [f]
        for i in range(1, n_max+1):
            f = f * i % mod
            fac.append(f)
        f = pow(f, mod-2, mod)
        self.facinv = facinv = [f]
        for i in range(n_max, 0, -1):
            f = f * i % mod
            facinv.append(f)
        facinv.reverse()

N = int(input())
G = [[] for _ in range(N+1)]
for _ in range(N-1):
    a, b = map(int, input().split())
    G[a].append(b)
    G[b].append(a)

Siz = [0] * (N+1)
A = [0] * (N+1)
mod = 10**9 + 7
comb = Combination(202020)
def dfs1(v, p):
    siz = 1
    an_prod = 1
    an_numer = 0
    an_denom = 1
    for u in G[v]:
        if u != p:
            s, a = dfs1(u, v)
            siz += s
            an_prod = an_prod * a % mod
            an_numer += s
            an_denom = an_denom * comb.facinv[s] % mod
    an_numer = comb.fac[an_numer]
    an = an_numer * an_denom % mod * an_prod % mod
    Siz[v] = siz
    A[v] = an
    return siz, an

_, ans = dfs1(1, 0)
# for i in range(1, N+1):
#     print(i, dfs1(i, 0))

Ans = [0] * (N+1)
Ans[1] = ans

def dfs2(v, p, an_p):
    siz_v = Siz[v]
    if v==1:
        ans = A[v]
    else:
        ans = A[v] * comb.fac[N-1] % mod * comb.facinv[siz_v-1] % mod * an_p % mod * comb.facinv[N-siz_v] % mod
        Ans[v] = ans
    for u in G[v]:
        if p != u:
            siz_u = Siz[u]
            an = ans * comb.facinv[N-1] % mod * comb.fac[N-1-siz_u] % mod * comb.fac[siz_u] % mod * pow(A[u], mod-2, mod) % mod
            dfs2(u, v, an)

dfs2(1, 0, None)
print("\n".join(map(str, Ans[1:])))

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
mod = 10**9+7

def pow_mod(base, exp):  # [numba_pow, "i8(i8,i8)"],
    # mod はグローバル変数を参照
    exp %= mod - 1
    res = 1
    while exp:
        if exp % 2:
            res = res * base % mod
        base = base * base % mod
        exp //= 2
    return res

def comb_cunstruct(n):
    # mod はグローバル変数を参照
    fac = np.empty(n+1, dtype=np.int64)
    facinv = np.empty(n+1, dtype=np.int64)
    fac[0] = f = 1
    for i in range(1, n+1):
        f = f * i % mod
        fac[i] = f
    f = pow_mod(f, -1)
    for i in range(n, -1, -1):
        facinv[i] = f
        f = f * i % mod
    return fac, facinv

def comb(n, r, fac, facinv):
    # mod はグローバル変数を参照
    return fac[n] * facinv[r] % mod * facinv[n-r] % mod

def rerooting(n, edges, fac, facinv):
    # 全方位木 dp
    # 参考: https://qiita.com/keymoon/items/2a52f1b0fb7ef67fb89e
    identity = (1, 0)
    def merge(a, b):
        return a[0] * b[0] % mod * comb(a[1] + b[1], a[1], fac, facinv) % mod, a[1] + b[1]
    def add_node(value, idx):
        return value[0], value[1] + 1
    merge(identity, identity)

    G = [[0]*0 for _ in range(n)]
    G_idxs = [[0]*0 for _ in range(n)]  # 自分を指すノードのインデックス
    for idx_edges in range(n-1):
        a, b = edges[idx_edges]
        G_idxs[a].append(len(G[b]))
        G_idxs[b].append(len(G[a]))
        G[a].append(b)
        G[b].append(a)
    # step 1
    parents = [0] * n
    order = []  # 行きがけ順
    stack = [0]
    parents[0] = -1
    while stack:
        v = stack.pop()
        p = parents[v]
        order.append(v)
        for u in G[v]:
            if p != u:
                stack.append(u)
                parents[u] = v
    subtree_res = [[identity]*len(Gv) for Gv in G]
    # 下から登る
    for v in order[:0:-1]:
        p = parents[v]
        result = identity
        for idx_Gv, (u, subtree_res_v_i) in enumerate(zip(G[v], subtree_res[v])):
            if p == u:
                parent_idx = idx_Gv
            else:
                result = merge(result, subtree_res_v_i)
        idx_p2v = G_idxs[v][parent_idx]
        subtree_res[p][idx_p2v] = add_node(result, v)
    # step 2
    # 上から降りる
    results = [identity] * n
    for v in order:
        subtree_res_v = subtree_res[v]
        cum = identity
        cum_from_tail = [identity]
        for r in subtree_res_v[:0:-1]:
            cum = merge(r, cum)
            cum_from_tail.append(cum)
        cum_from_tail.reverse()
        cum = identity
        for r, cum_t, u, idx_u2v in zip(subtree_res_v, cum_from_tail, G[v], G_idxs[v]):
            result = add_node(merge(cum, cum_t), v)
            subtree_res[u][idx_u2v] = result
            cum = merge(cum, r)
        results[v] = add_node(cum, v)
    return np.array(results)


numba_compile([
    [pow_mod, "i8(i8,i8)"],
    [comb_cunstruct, "Tuple((i8[:],i8[:]))(i8,)"],
    [comb, "i8(i8,i8,i8[:],i8[:])"],
    [rerooting, "(i8,i8[:,:],i8[:],i8[:])"],
])


def main():
    N = int(input())
    AB = np.array(sys.stdin.read().split(), dtype=np.int64).reshape(N-1, 2) - 1
    fac, facinv = comb_cunstruct(202020)
    Ans = rerooting(N, AB, fac, facinv)
    print("\n".join(map(str, Ans[:, 0].tolist())))

main()

Buggy submission (verdict: Runtime Error):
import numpy as np
mod = 10**9+7

def pow_mod(base, exp):  # [numba_pow, "i8(i8,i8)"],
    # mod はグローバル変数を参照
    exp %= mod - 1
    res = 1
    while exp:
        if exp % 2:
            res = res * base % mod
        base = base * base % mod
        exp //= 2
    return res

def comb_cunstruct(n):
    # mod はグローバル変数を参照
    fac = np.empty(n+1, dtype=np.int64)
    facinv = np.empty(n+1, dtype=np.int64)
    fac[0] = f = 1
    for i in range(1, n+1):
        f = f * i % mod
        fac[i] = f
    f = pow_mod(f, -1)
    for i in range(n, -1, -1):
        facinv[i] = f
        f = f * i % mod
    return fac, facinv

def comb(n, r, fac, facinv):
    # mod はグローバル変数を参照
    return fac[n] * facinv[r] % mod * facinv[n-r] % mod

def rerooting(n, edges, fac, facinv):
    # 全方位木 dp
    # 参考: https://qiita.com/keymoon/items/2a52f1b0fb7ef67fb89e
    identity = (1, 0)
    def merge(a, b):
        return a[0] * b[0] % mod * comb(a[1] + b[1], a[1], fac, facinv) % mod, a[1] + b[1]
    def add_node(value, idx):
        return value[0], value[1] + 1
    merge(identity, identity)

    G = [[0]*0 for _ in range(n)]
    G_idxs = [[0]*0 for _ in range(n)]  # 自分を指すノードのインデックス
    for idx_edges in range(n-1):
        a, b = edges[idx_edges]
        G_idxs[a].append(len(G[b]))
        G_idxs[b].append(len(G[a]))
        G[a].append(b)
        G[b].append(a)
    # step 1
    parents = [0] * n
    order = []  # 行きがけ順
    stack = [0]
    parents[0] = -1
    while stack:
        v = stack.pop()
        p = parents[v]
        order.append(v)
        for u in G[v]:
            if p != u:
                stack.append(u)
                parents[u] = v
    subtree_res = [[identity]*len(Gv) for Gv in G]
    # 下から登る
    for v in order[:0:-1]:
        p = parents[v]
        result = identity
        for idx_Gv, (u, subtree_res_v_i) in enumerate(zip(G[v], subtree_res[v])):
            if p == u:
                parent_idx = idx_Gv
            else:
                result = merge(result, subtree_res_v_i)
        idx_p2v = G_idxs[v][parent_idx]
        subtree_res[p][idx_p2v] = add_node(result, v)
    # step 2
    # 上から降りる
    results = [identity] * n
    for v in order:
        subtree_res_v = subtree_res[v]
        cum = identity
        cum_from_tail = [identity]
        for r in subtree_res_v[:0:-1]:
            cum = merge(r, cum)
            cum_from_tail.append(cum)
        cum_from_tail.reverse()
        cum = identity
        for r, cum_t, u, idx_u2v in zip(subtree_res_v, cum_from_tail, G[v], G_idxs[v]):
            result = add_node(merge(cum, cum_t), v)
            subtree_res[u][idx_u2v] = result
            cum = merge(cum, r)
        results[v] = add_node(cum, v)
    return np.array(results)


numba_compile([
    [pow_mod, "i8(i8,i8)"],
    [comb_cunstruct, "Tuple((i8[:],i8[:]))(i8,)"],
    [comb, "i8(i8,i8,i8[:],i8[:])"],
    [rerooting, "(i8,i8[:,:],i8[:],i8[:])"],
])


def main():
    N = int(input())
    AB = np.array(sys.stdin.read().split(), dtype=np.int64).reshape(N-1, 2) - 1
    fac, facinv = comb_cunstruct(202020)
    Ans = rerooting(N, AB, fac, facinv)
    print("\n".join(map(str, Ans[:, 0].tolist())))

main()