Problem p03181 (python)

Fixed submission:
# 参考 -> https://atcoder.jp/contests/dp/submissions/8271798

import itertools
import sys
input = sys.stdin.readline

N, MOD = map(int, input().split())
XY = [list(map(int, input().split())) for i in range(N-1)]
prod = lambda x, y: x*y%MOD

tree = [[] for i in range(N)]
for x, y in XY:
  tree[x-1].append(y-1)
  tree[y-1].append(x-1)

root = 0
parent = [-1] * (N+1)
order = []
stack = [root]
while stack:
  node = stack.pop()
  order.append(node)
  for child in tree[node]:
    if child == parent[node]:
      continue
    parent[child] = node
    stack.append(child)

# dp_down[i] : 頂点iを黒く塗ったときの子側の塗り方
dp_down = [0] * N
for node in order[::-1]: # 葉から根に向かって伝搬させる
  res = 1
  for child in tree[node]:
    res *= dp_down[child] + 1
    res %= MOD
  dp_down[node] = res

# dp_up[i] : 頂点iを黒く塗ったときの親側の塗り方
dp_up = [1] * N
for node in order: # 根から葉に向かって伝搬させる
  par = parent[node]
  arr = [dp_down[c]+1 if c!=par else dp_up[node] for c in tree[node]]
  # それぞれの子(=c)について、自身を根とする下側の部分木を除いた積
  left = [1] + list(itertools.accumulate(arr, prod))[:-1]
  right = list(itertools.accumulate(reversed(arr), prod))[-2::-1] + [1]
  upper_cases = [l*r%MOD for l, r in zip(left, right)]
  for c, res in zip(tree[node], upper_cases):
    if c != par:
      dp_up[c] = res + 1 # (+1 : cより上側が全て白である場合を足し合わせる)

ans = [d*u%MOD for d, u in zip(dp_down, dp_up)]
print(*ans, sep='\n')

Buggy submission (verdict: Runtime Error):
def power(n, r, mod=10**9+7):
    if r == 0: return 1
    if r%2 == 0:
        return power(n*n % mod, r//2, mod) % mod
    if r%2 == 1:
        return n * power(n, r-1, mod) % mod


N, M = map(int, input().split())
tree = [[] for i in range(N+1)]
for i in range(N-1):
  x, y = map(int, input().split())
  tree[x].append(y)
  tree[y].append(x)

dp1 = [0] * (N+1)
dp2 = [0] * (N+1)
parents = [-1] * (N+1)
leaf_nodes = []

def solve1(node, parent):
  global dp1
  global parents
  global leaf_nodes
  parents[node] = parent
  if dp1[node] > 0:
    return dp1[node]
  else:
    children = tree[node]
    t = 1
    for child in children:
      if child == parent:
        continue
      
      t *= (solve1(child, node) + 1)
      t %= M
    
    if t == 1:
      leaf_nodes.append(node)
    
    dp1[node] = t
    return t


def solve2(node):
  global dp2
  global parents
  parent = parents[node]
  if dp2[node] > 0:
    return dp2[node]
  else:
    children = tree[node]
    t = 1
    for child in children:
      if child == parent:
        t *= (solve2(parent)//(dp1[node] + 1) + 1)
      else:
        t *= (dp1[child] + 1)
    
    dp2[node] = t
    return t
  

_ = solve1(1, 0)
for leaf in leaf_nodes:
  _ = solve2(leaf)


for ans in dp2[1:]:
  print(ans%M)

Buggy submission (verdict: Runtime Error):
N, M = map(int, input().split())

tree = [[] for i in range(N+1)]
for i in range(N-1):
  x, y = map(int, input().split())
  tree[x].append(y)
  tree[y].append(x)

dp1 = [0] * (N+1)
dp2 = [0] * (N+1)
parents = [-1] * (N+1)
leaf_nodes = []

def solve1(node, parent):
  global dp1
  global parents
  global leaf_nodes
  parents[node] = parent
  if dp1[node] > 0:
    return dp1[node]
  else:
    children = tree[node]
    t = 1
    for child in children:
      if child == parent:
        continue
      
      t *= (solve1(child, node) + 1)
      t %= M
    
    if t == 1:
      leaf_nodes.append(node)
    
    dp1[node] = t
    return t


def solve2(node):
  global dp2
  global parents
  parent = parents[node]
  if dp2[node] > 0:
    return dp2[node]
  else:
    children = tree[node]
    t = 1
    for child in children:
      if child == parent:
        t *= (solve2(parent)//(dp1[node] + 1) + 1)
        t %= M
      else:
        t *= (dp1[child] + 1)
        t %= M
    
    dp2[node] = t
    return t
  

_ = solve1(1, 0)
for leaf in leaf_nodes:
  _ = solve2(leaf)

print(*dp2[1:], sep='\n')