Problem p03181 (Python) — buggy submission, verdict: Time Limit Exceeded

# これは逆元を用いてやる？
# この問題の場合、逆元をうまく設定するのが難しいので
# 左右からのmergeになりそう
from collections import deque
import sys

input = sys.stdin.buffer.readline


n, m = map(int, input().split())
# 各々の頂点の次数
deg = [0] * (n + 1)
# 隣接リストにグラフをまとめる
graph = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    x, y = map(int, input().split())
    graph[x].append(y)
    graph[y].append(x)
    deg[x] += 1
    deg[y] += 1

# dp1は葉向きに見たときの塗り方
# 自身が黒い時、下の候補はどうなるか？
dp1 = [0] * (n + 1)

# dp2は根向きに見たときの塗り方
# 自身が黒い時、親側の候補はどうなるか？
dp2 = [0] * (n + 1)

# 1回目
start = 1
par = [-1] * (n + 1)
que = deque([start])
topo = []
while que:
    v = que.popleft()
    topo.append(v)
    for i in range(deg[v]):
        u = graph[v][i]
        if par[v] == u:
            continue
        par[u] = v
        que.append(u)

# dpするついでにi=0からの積とi=deg[v]-1からの積を取っておくか。
dpl = [[] for _ in range(n + 1)]
dpr = [[] for _ in range(n + 1)]

# 頂点番号から何番目かがわかるようにしたいから辞書で持っておく？
edge_num = [{} for _ in range(n + 1)]

for v in topo[::-1]:
    temp = 1
    for i in range(deg[v]):
        u = graph[v][i]
        edge_num[v][u] = i  # edge_num[v]で頂点vに関する辞書。頂点uからidxが欲しいので記録
        if u == par[v]:
            dpl[v].append(temp)
            continue
        temp *= dp1[u] + 1
        temp %= m
        dpl[v].append(temp)
    dp1[v] = temp
    temp = 1
    for i in range(deg[v] - 1, -1, -1):
        u = graph[v][i]
        if u == par[v]:
            dpr[v].append(temp)
            continue
        temp *= dp1[u] + 1
        temp %= m
        dpr[v].append(temp)
# これで普通の木DPに加えて
# 左右からの累積積
# 頂点番号からidxへの辞書
# を獲得したはず

# dpは欲しい答えとしておく
# dp1とdp2から求める
dp = [0] * (n + 1)

for v in topo:
    # dp2の計算
    p = par[v]
    if p == -1:  # 例外、親がいない
        dp[v] = dp1[v]
        dp[v] %= m
        continue
    # 根側の処理
    # 自分以外の積を作りたい
    idx = edge_num[p][v]
    num = deg[p]
    dp2[v] = 1
    if idx - 1 >= 0:
        dp2[v] *= dpl[p][idx - 1]
        dp2[v] %= m
    if num - 2 - idx >= 0:
        dp2[v] *= dpr[p][num - 2 - idx]
        dp2[v] %= m
    # 根側の根側の処理
    pp = par[p]  # 親の親
    if pp != -1:  # 親の親がいるならば
        dp2[v] *= dp2[p] + 1  # 白=1か黒=dp2[pp]か
        dp2[v] %= m
    # dpの計算
    dp[v] = dp1[v] * (dp2[v] + 1)
    dp[v] %= m
print(*dp[1:], sep="\\n")


Error: TLE

Fixed submission:

# これは逆元を用いてやる方法
from collections import deque
import sys

input = sys.stdin.buffer.readline


n, m = map(int, input().split())
graph = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    x, y = map(int, input().split())
    graph[x].append(y)
    graph[y].append(x)

# dp1は葉向きに見たときの塗り方
# 自身が黒い時、下の候補はどうなるか？
dp1 = [0] * (n + 1)
# dp2は根向きに見たときの塗り方
# 自身が黒い時、親側の候補はどうなるか？
dp2 = [0] * (n + 1)

# 1回目
start = 1
par = [-1] * (n + 1)
que = deque([start])
topo = []
while que:
    v = que.popleft()
    topo.append(v)
    for u in graph[v]:
        if par[v] == u:
            continue
        par[u] = v
        que.append(u)

for v in topo[::-1]:
    dp1[v] = 1
    for u in graph[v]:
        if u == par[v]:
            continue
        dp1[v] *= dp1[u] + 1

# dpは欲しい答えとしておく
# dp1とdp2から求める
dp = [0] * (n + 1)

for v in topo:
    # dp2の計算
    p = par[v]
    if p == -1:
        dp[v] = dp1[v]
        dp[v] %= m
        continue
    dp2[v] = dp1[p]
    dp2[v] //= dp1[v] + 1  # 親側に関して自身の関与は消す
    pp = par[p]  # 親の親
    if pp != -1:  # 親の親がいるならば
        dp2[v] *= dp2[p] + 1  # 白=1か黒=dp2[pp]か

    # dpの計算
    dp[v] = dp1[v] * (dp2[v] + 1)
    dp[v] %= m
print(*dp[1:], sep="\\n")