Problem p02728 (python)

Fixed submission:

import sys
from functools import reduce
def rerooting(N, adj, merge, finalize,identity):
    """
    merge: (T,T) -> T
    This merges two dp results. (T, operator) must be a monoid.

    identity: operator(x, identity) = operator(identity, x) = x
    
    finalize: T -> T
    This transforms the merged result to the dp result
    """

    order = [None]*N
    parent = [None]*N
    parent[0] = -1
    stack = [0]
    for i in range(N):
        v = stack.pop()
        order[i] = v
        for u in adj[v]:
            if parent[u] is None:
                parent[u] = v
                stack.append(u)

    # from leaf to parent
    dp_down = [None]*N
    for v in reversed(order[1:]):
        p = parent[v]
        dp_down[v] = finalize(reduce(merge,
            (dp_down[c] for c in adj[v] if c != p),
            identity))

    # from parent to leaf
    dp_up = [None]*N
    dp_up[0] = identity
    for v in order:
        if len(adj[v]) == 1 and adj[v][0] == parent[v]:
            continue
        p = parent[v]
        left = [dp_up[v]]
        for u in adj[v]:
            if u != p:
                left.append(merge(left[-1],dp_down[u]))

        right = [identity]
        for u in reversed(adj[v]):
            if u != p:
                right.append(merge(dp_down[u],right[-1]))
        i = 0
        for u in adj[v]:
            if u != p:
                dp_up[u] = finalize(merge(left[i],right[-i-2]))
                i += 1

    for v,l in enumerate(adj):
        order[v] = reduce(merge,
                    (dp_down[u] for u in adj[v] if u != parent[v]),
                    identity)
        order[v] = finalize(merge(order[v], dp_up[v]))

    return order


def solve(N,edges):
    MOD = 10**9+7

    factorio = [1]*(N+1)
    t = 1
    for i in range(1,N+1):
        t *= i
        t %= MOD
        factorio[i] = t

    adj = [[] for _ in range(N)]
    for a,b in edges:
        adj[a].append(b)
        adj[b].append(a)

    def func(x,y):
        cx,nx = x
        cy,ny = y
        c = cx*cy*factorio[nx+ny]*pow(factorio[nx]*factorio[ny],MOD-2,MOD)
        c %= MOD
        n = nx+ny
        return c, n

    def func2(x):
        c,n = x
        return c,n+1

    identity = (1,0)

    res = rerooting(N, adj, func, func2, identity)

    for k,n in res:
        print(k)


if __name__ == '__main__':
    read = sys.stdin.buffer.read
    readline = sys.stdin.buffer.readline
    readlines = sys.stdin.buffer.readlines
     
    N = int(readline())
    m = map(int, read().split())
    edges = tuple((a-1,b-1) for a,b in zip(m, m))

    solve(N, edges)

Buggy submission (verdict: Runtime Error):
from itertools import accumulate,chain
import sys

sys.setrecursionlimit(10000000)

MOD = 10**9+7

# factorio = tuple(chain((1,),accumulate(range(1,N+1), func=lambda x,y: (x*y)%MOD)))

from itertools import chain
def rerooting(N, adj, func):
    """
    func: iterable<T> -> T
    func takes all subtree dp results and 
    returns its aggregation result.
    """

    stack = [0]
    order = [None]*N
    parent = [None]*N
    parent[0] = -1
    for i in range(N):
        v = stack.pop()
        order[i] = v
        for u in adj[v]:
            if parent[u] is None:
                parent[u] = v
                stack.append(u)

    # from leaf to parent
    dp_down = [None]*N
    for v in reversed(order[1:]):
        dp_down[v] = func((dp_down[c] for c in adj[v] if c != parent[v]))

    dp_up = [None]*N
    for v in order[1:]:
        p = parent[v]
        dp_up[v] = func(
            chain(
                (dp_down[u] for u in adj[p] if u not in (v,parent[p])),
                (dp_up[p],) if p != 0 else ()))

    for v,l in enumerate(adj):
        order[v] = func(
            chain(
                (dp_down[u] for u in adj[v] if u != parent[v]),
                (dp_up[v],) if v != 0 else ()))

    return order


def solve(N,edges):
    factorio = [1]*(N+1)
    t = 1
    for i in range(1,N+1):
        t *= i
        t %= MOD
        factorio[i] = t

    adj = [[] for _ in range(N)]
    for a,b in edges:
        adj[a].append(b)
        adj[b].append(a)

    def aggregator(it):
        comb = 1
        n_children = 0
        for c, n in it:
            comb *= pow(factorio[n],-1,MOD)*c
            comb %= MOD
            n_children += n
        comb *= factorio[n_children]
        comb %= MOD
        return comb, n_children+1

    res = rerooting(N,adj,aggregator)

    for k,n in res:
        print(k)


if __name__ == '__main__':
    read = sys.stdin.buffer.read
    readline = sys.stdin.buffer.readline
    readlines = sys.stdin.buffer.readlines
     
    N = int(readline())
    m = map(int, read().split())
    edges = tuple((a-1,b-1) for a,b in zip(m, m))

    solve(N, edges)

Buggy submission (verdict: Runtime Error):
from itertools import accumulate,chain
from sys import setrecursionlimit

setrecursionlimit(10000000)

MOD = 10**9+7


# 拡張ユークリッド互除法
# ax + by = gcd(a,b)の最小整数解を返す
def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)


# mを法とするaの乗法的逆元
def modinv(a, m):
    g, x, y = egcd(a, m)
    if g != 1:
        raise Exception('modular inverse does not exist')
    else:
        return x % m




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

factorio = tuple(chain((1,),accumulate(range(1,N+1), lambda x,y: (x*y)%MOD)))

patterns = dict()
childrens = dict()
def rec(a,b):
    if (a,b) in patterns:
        return patterns[(a,b)], childrens[(a,b)]
    if a != -1 and len(adj[b]) == 1:
        patterns[(a,b)] = 1
        childrens[(a,b)] = 1
        return 1,1

    prod = 1 # product of patterns
    d = 1 # denominator of multinomial coefficient
    n = 0 # total childrens
    for c in adj[b]:
        if a == c:
            continue
        p,c = rec(b,c)
        prod *= p
        prod %= MOD
        d *= factorio[c]
        d %= MOD
        n += c

    patterns[(a,b)] = (prod*factorio[n]*modinv(d,MOD))%MOD
    childrens[(a,b)] = n+1
    return patterns[(a,b)], childrens[(a,b)]

for a,b in edges:
    rec(a,b)
    rec(b,a)

for k in range(N):
    p,c = rec(-1,k)
    print(p)

Buggy submission (verdict: Runtime Error):
from itertools import accumulate,chain

MOD = 10**9+7


# 拡張ユークリッド互除法
# ax + by = gcd(a,b)の最小整数解を返す
def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)


# mを法とするaの乗法的逆元
def modinv(a, m):
    g, x, y = egcd(a, m)
    if g != 1:
        raise Exception('modular inverse does not exist')
    else:
        return x % m




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

factorio = tuple(chain((1,),accumulate(range(1,N+1), lambda x,y: (x*y)%MOD)))

patterns = dict()
childrens = dict()
def rec(a,b):
    if (a,b) in patterns:
        return patterns[(a,b)], childrens[(a,b)]
    if a != -1 and len(adj[b]) == 1:
        patterns[(a,b)] = 1
        childrens[(a,b)] = 1
        return 1,1

    prod = 1 # product of patterns
    d = 1 # denominator of multinomial coefficient
    n = 0 # total childrens
    for c in adj[b]:
        if a == c:
            continue
        p,c = rec(b,c)
        prod *= p
        prod %= MOD
        d *= factorio[c]
        d %= MOD
        n += c

    patterns[(a,b)] = (prod*factorio[n]*modinv(d,MOD))%MOD
    childrens[(a,b)] = n+1
    return patterns[(a,b)], childrens[(a,b)]

for a,b in edges:
    rec(a,b)
    rec(b,a)

for k in range(N):
    p,c = rec(-1,k)
    print(p)