Problem p02985 (python)

Fixed submission:
import os
import sys
from collections import deque
from functools import lru_cache, reduce
from operator import mul

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(2147483647)
INF = float("inf")
IINF = 10 ** 18
MOD = 1000000007


def get_factorials(max, mod=None):
    """
    階乗 0!, 1!, 2!, ..., max!
    :param int max:
    :param int mod:
    :return:
    """
    ret = [1]
    n = 1
    if mod:
        for i in range(1, max + 1):
            n *= i
            n %= mod
            ret.append(n)
    else:
        for i in range(1, max + 1):
            n *= i
            ret.append(n)
    return ret


def mod_invs(max, mod):
    """
    :param max:
    :param mod:
    :return:
    """
    invs = [1] * (max + 1)
    for x in range(2, max + 1):
        invs[x] = (-(mod // x) * invs[mod % x]) % mod
    return invs


N, K = list(map(int, sys.stdin.readline().split()))
AB = [list(map(int, sys.stdin.readline().split())) for _ in range(N - 1)]

factorials = get_factorials(max=K, mod=MOD)
invs = mod_invs(K, MOD)


def mod_inv(a, mod):
    """
    a の逆元
    :param int a:
    :param int mod:
    :return:
    """
    return pow(a, mod - 2, mod)


@lru_cache(maxsize=None)
def npr(n, r, mod=None):
    """
    :param int n:
    :param int r:
    :param int mod: 3 以上の素数であること
    :rtype: int
    """
    if n < r:
        return 0

    # 何度も呼ぶ場合は最大の n 以下の階乗を事前に計算しておくといい
    # return factorials[n] * mod_inv(factorials[r], mod) * mod_inv(factorials[n - r], mod) % mod
    if r == 0:
        return 1

    return factorials[n] * mod_inv(factorials[n-r], mod) % MOD


if N == 1:
    print(K)
    exit()

graph = [[] for _ in range(N + 1)]
for a, b in AB:
    graph[a].append(b)
    graph[b].append(a)
# 葉っぱを探す
leaf = None
for i in range(N):
    if len(graph[i]) == 1:
        leaf = i
        break

v = graph[leaf][0]
ans = npr(K, len(graph[v]) + 1, MOD)
seen = [False] * (N + 1)
seen[v] = True
que = deque()
for u in graph[v]:
    que.append(u)
    seen[u] = True

while que:
    v = que.popleft()
    seen[v] = True
    if len(graph[v]) == 1:
        continue
    ans *= npr(K - 2, len(graph[v]) - 1, MOD)
    ans %= MOD
    for u in graph[v]:
        if not seen[u]:
            que.append(u)

print(ans)

Buggy submission (verdict: Runtime Error):
import os
import sys
from collections import deque
from functools import lru_cache, reduce
from operator import mul

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(2147483647)
INF = float("inf")
IINF = 10 ** 18
MOD = 1000000007



@lru_cache(maxsize=None)
def npr(n, r, mod=None):
    """
    :param int n:
    :param int r:
    :param int mod: 3 以上の素数であること
    :rtype: int
    """
    if n < r:
        return 0

    # 何度も呼ぶ場合は最大の n 以下の階乗を事前に計算しておくといい
    # return factorials[n] * mod_inv(factorials[r], mod) * mod_inv(factorials[n - r], mod) % mod
    if r == 0:
        return 1
    if mod:
        return reduce(mul, range(n, n - r, -1)) % mod
    else:
        # math.factorial よりこっちのが速い
        # https://atcoder.jp/contests/abc110/submissions?f.Task=&f.Language=&f.Status=&f.User=nohtaray
        return reduce(mul, range(n, n - r, -1))


N, K = list(map(int, sys.stdin.readline().split()))
AB = [list(map(int, sys.stdin.readline().split())) for _ in range(N - 1)]

graph = [[] for _ in range(N + 1)]
for a, b in AB:
    graph[a].append(b)
    graph[b].append(a)
# 葉っぱを探す
leaf = None
for i in range(N):
    if len(graph[i]) == 1:
        leaf = i
        break


v = graph[leaf][0]
ans = npr(K, len(graph[v]) + 1, MOD)
seen = [False] * (N + 1)
seen[v] = True
que = deque()
for u in graph[v]:
    que.append(u)
    seen[u] = True

while que:
    v = que.popleft()
    seen[v] = True
    if len(graph[v]) == 1:
        continue
    ans *= npr(K - 2, len(graph[v]) - 1, MOD)
    ans %= MOD
    for u in graph[v]:
        if not seen[u]:
            que.append(u)

print(ans)

Buggy submission (verdict: Runtime Error):
import os
import sys
from collections import deque
from functools import lru_cache, reduce
from operator import mul

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(2147483647)
INF = float("inf")
IINF = 10 ** 18
MOD = 1000000007



@lru_cache(maxsize=None)
def npr(n, r, mod=None):
    """
    :param int n:
    :param int r:
    :param int mod: 3 以上の素数であること
    :rtype: int
    """
    if n < r:
        return 0

    # 何度も呼ぶ場合は最大の n 以下の階乗を事前に計算しておくといい
    # return factorials[n] * mod_inv(factorials[r], mod) * mod_inv(factorials[n - r], mod) % mod
    if r == 0:
        return 1
    if mod:
        return reduce(mul, range(n, n - r, -1)) % mod
    else:
        # math.factorial よりこっちのが速い
        # https://atcoder.jp/contests/abc110/submissions?f.Task=&f.Language=&f.Status=&f.User=nohtaray
        return reduce(mul, range(n, n - r, -1))


N, K = list(map(int, sys.stdin.readline().split()))
AB = [list(map(int, sys.stdin.readline().split())) for _ in range(N - 1)]

graph = [[] for _ in range(N + 1)]
for a, b in AB:
    graph[a].append(b)
    graph[b].append(a)
# 葉っぱを探す
leaf = None
for i in range(N):
    if len(graph[i]) == 1:
        leaf = i
        break


v = graph[leaf][0]
ans = npr(K, len(graph[v]) + 1, MOD)
seen = [False] * (N + 1)
seen[v] = True
que = deque()
for u in graph[v]:
    que.append(u)
    seen[u] = True

while que:
    v = que.popleft()
    seen[v] = True
    if len(graph[v]) == 1:
        continue
    ans *= npr(K - 2, len(graph[v]) - 1, MOD)
    ans %= MOD
    for u in graph[v]:
        if not seen[u]:
            que.append(u)

print(ans)