Problem p02728 (python)

Fixed submission:
from functools import *
from collections import *
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.buffer.readline
 
M = 10**9+7

@lru_cache(maxsize=None)
def mod_inv(a):
  b, x, y = M, 1, 0
  while a != 1:
    (d, a), b = divmod(b, a), a
    x, y = y-d*x, x
  return x%M

@lru_cache(maxsize=None)
def calc_subtree(v):
  ret = 1
  size = 1
  for child in Edge[v]:
    Edge[child].remove(v)
    r, s = calc_subtree(child)
    ret = ret * r % M
    size += s
      
  ret = ret * size % M
  return ret, size

ans = {}
def set_ans(v, a):
  ans[v] = a
  for child in Edge[v]:
    n = calc_subtree(child)[1]
    set_ans(child, ans[v] * n * mod_inv(N-n) % M)

N = int(input())
Edge = defaultdict(set)
fact = N
for i in range(1, N):
  fact = fact * i % M
  a, b = map(int, input().split())
  Edge[a].add(b)
  Edge[b].add(a)

set_ans(1, mod_inv(calc_subtree(1)[0]) * fact % M)
for i in range(1, N+1):
  print(ans[i])

Buggy submission (verdict: Runtime Error):
from functools import *
from collections import *
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.buffer.readline
 
M = 10**9+7
N = int(input())

@lru_cache(maxsize=None)
def mod_inv(x):
  if x == 1:
    return 1
  return M // x * -mod_inv(M%x) % M

weght = [0]*(N+1)
size = [0]*(N+1)
def calc_subtree(v):
  W = 1
  S = 1
  for child in Edge[v]:
    Edge[child].remove(v)
    w, s = calc_subtree(child)
    W = W * w % M
    S += s
  weight[v] = W * S % M
  size[v] = S
  return weight[v], size[v]

ans = [0]*(N+1)
def set_ans(v, a):
  ans[v] = a
  for child in Edge[v]:
    n = calc_subtree(child)[1]
    set_ans(child, ans[v] * n * mod_inv(N-n) % M)

Edge = defaultdict(set)
fact = N
for i in range(1, N):
  fact = fact * i % M
  a, b = map(int, input().split())
  Edge[a].add(b)
  Edge[b].add(a)

set_ans(1, mod_inv(calc_subtree(1)[0]) * fact % M)
for i in range(1, N+1):
  print(ans[i])

Buggy submission (verdict: Runtime Error):
from functools import *
from collections import *

M = 10**9+7

@lru_cache(maxsize=None)
def mod_inv(a):
  b, x, y = M, 1, 0
  while a != 1:
    (d, a), b = divmod(b, a), a
    x, y = y-d*x, x
  return x%M

cache = defaultdict(dict)
def calc_cache(self, parent):
  c = cache[self]
  if not c or parent in c:
    return c[parent]
  
  if c and None not in c:
    for p, (r2, t2, s2) in c.items():
      r3, t3, s3 = calc_cache(p, self)
      r = r2 * t3
      s = s2 + s3
      c[None] = r, None, s
      break
  
  if parent is None:
    return c[parent]
  
  if None in c:
    r0, t0, s0 = c[None]
    r1, t1, s1 = calc_cache(parent, self)
    r = r0 * mod_inv(t1) % M
    s = s0 - s1
    c[parent] = r, r * mod_inv(s) % M, s
    return c[parent]

N = int(input())
Edge = defaultdict(set)
fact = 1
for i in range(1, N):
  fact = fact * i % M
  a, b = map(int, input().split())
  Edge[a].add(b)
  Edge[b].add(a)

cache = defaultdict(dict)
stack = [None, 1]
children = {}
while len(stack) > 1:
  parent = stack[-2]
  self = stack[-1]
  if self not in children:
    children[self] = Edge[self]-{parent}
  if children[self]:
    stack.append(children[self].pop())
  else:
    ret = 1
    size = 1
    for child in (Edge[self] - {parent}):
      r, t, s = cache[child][self]
      ret = ret * t % M
      size += s
    cache[self][parent] = ret, ret * mod_inv(size) % M if parent else None, size
    stack.pop()

for i in range(1, N+1):
  print(calc_cache(i, None)[0] * fact % M)