Problem p02985 (python)

Fixed submission:
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)

N,K = map(int,input().split())
MOD = 10 ** 9 + 7

edge = [[] for _ in range(N)]
for _ in range(N-1):
  a,b = map(lambda x: int(x) - 1,input().split())
  edge[a].append(b)
  edge[b].append(a)
deg = [len(E) for E in edge]
  
U = 10 ** 5 + 100
fact = [1] * (U+1)
fact_inv = [1] * (U+1)
for n in range(1,U+1):
  fact[n] = fact[n-1] * n % MOD
fact_inv[U] = pow(fact[U],MOD-2,MOD)
for n in range(U-1,0,-1):
  fact_inv[n] = fact_inv[n+1] * (n+1) % MOD

def perm(n,k):
  if k > n or k < 0:
    return 0
  return fact[n] * fact_inv[n-k] % MOD

def F(x,parent = None):
  if parent is None:
    c = perm(K, deg[x] + 1)
  else:
    c = perm(K-2, deg[x] - 1)
  for y in edge[x]:
    if y == parent:
      continue
    c *= F(y,x)
    c %= MOD
  return c

answer = F(0)
print(answer)

Buggy submission (verdict: Runtime Error):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

import numpy as np

MOD = 10**9 + 7

N,K = map(int,readline().split())
AB = map(int,read().split())

if K == 1:
    if N == 1:
        print(1)
    else:
        print(0)
    raise Exception
    exit()

def cumprod(arr,MOD):
    L = len(arr); Lsq = int(L**.5+1)
    arr = np.resize(arr,Lsq**2).reshape(Lsq,Lsq)
    for n in range(1,Lsq):
        arr[:,n] *= arr[:,n-1]; arr[:,n] %= MOD
    for n in range(1,Lsq):
        arr[n] *= arr[n-1,-1]; arr[n] %= MOD
    return arr.ravel()[:L]

def make_fact(U,MOD):
    x = np.arange(U,dtype=np.int64); x[0] = 1
    fact = cumprod(x,MOD)
    x = np.arange(U,0,-1,dtype=np.int64); x[0] = pow(int(fact[-1]),MOD-2,MOD)
    fact_inv = cumprod(x,MOD)[::-1]
    return fact,fact_inv

fact,fact_inv = make_fact(K+10,MOD)

graph = [[] for _ in range(N+1)]
for a,b in zip(AB,AB):
    graph[a].append(b)
    graph[b].append(a)

# 次数1の頂点を根とする
deg = [len(x) for x in graph]
for i,x in enumerate(deg):
    if x == 1:
        break
root = i

parent = [0] * (N+1)
ts = [root]
stack = [root]
while stack:
    x = stack.pop()
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x
        stack.append(y)
        ts.append(y)

perm = np.zeros(N+K+10,np.int64)
perm[:K-1] = fact[K-2] * fact_inv[K-2::-1] % MOD
perm = perm.tolist()

dp = [1] * (N+1)
for v in ts[::-1]:
    if v == root:
        break
    d = deg[v]
    dp[v] *= perm[d-1]; dp[v] %= MOD
    p = parent[v]
    dp[p] *= dp[v]; dp[p] %= MOD

x = dp[root]
x *= K*(K-1)%MOD
x %= MOD
print(x)

Buggy submission (verdict: Runtime Error):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

import numpy as np

MOD = 10**9 + 7

N,K = map(int,readline().split())
AB = map(int,read().split())

def cumprod(arr,MOD):
    L = len(arr); Lsq = int(L**.5+1)
    arr = np.resize(arr,Lsq**2).reshape(Lsq,Lsq)
    for n in range(1,Lsq):
        arr[:,n] *= arr[:,n-1]; arr[:,n] %= MOD
    for n in range(1,Lsq):
        arr[n] *= arr[n-1,-1]; arr[n] %= MOD
    return arr.ravel()[:L]

def make_fact(U,MOD):
    x = np.arange(U,dtype=np.int64); x[0] = 1
    fact = cumprod(x,MOD)
    x = np.arange(U,0,-1,dtype=np.int64); x[0] = pow(int(fact[-1]),MOD-2,MOD)
    fact_inv = cumprod(x,MOD)[::-1]
    return fact,fact_inv

fact,fact_inv = make_fact(K+10,MOD)

graph = [[] for _ in range(N+1)]
for a,b in zip(AB,AB):
    graph[a].append(b)
    graph[b].append(a)

# 次数1の頂点を根とする
deg = [len(x) for x in graph]
for i,x in enumerate(deg):
    if x == 1:
        break
root = i

parent = [0] * (N+1)
ts = [root]
stack = [root]
while stack:
    x = stack.pop()
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x
        stack.append(y)
        ts.append(y)

perm = np.zeros(N+K+10,np.int64)
perm[:K-1] = fact[K-2] * fact_inv[K-2::-1] % MOD
perm = perm.tolist()

dp = [1] * (N+1)
for v in ts[::-1]:
    if v == root:
        break
    d = deg[v]
    dp[v] *= perm[d-1]; dp[v] %= MOD
    p = parent[v]
    dp[p] *= dp[v]; dp[p] %= MOD

x = dp[root]
x *= K*(K-1)%MOD
x %= MOD
print(x)

Buggy submission (verdict: Runtime Error):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

import numpy as np

MOD = 10**9 + 7

N,K = map(int,readline().split())
AB = map(int,read().split())

def cumprod(arr,MOD):
    L = len(arr); Lsq = int(L**.5+1)
    arr = np.resize(arr,Lsq**2).reshape(Lsq,Lsq)
    for n in range(1,Lsq):
        arr[:,n] *= arr[:,n-1]; arr[:,n] %= MOD
    for n in range(1,Lsq):
        arr[n] *= arr[n-1,-1]; arr[n] %= MOD
    return arr.ravel()[:L]

def make_fact(U,MOD):
    x = np.arange(U,dtype=np.int64); x[0] = 1
    fact = cumprod(x,MOD)
    x = np.arange(U,0,-1,dtype=np.int64); x[0] = pow(int(fact[-1]),MOD-2,MOD)
    fact_inv = cumprod(x,MOD)[::-1]
    return fact,fact_inv

fact,fact_inv = make_fact(K+10,MOD)

graph = [[] for _ in range(N+1)]
for a,b in zip(AB,AB):
    graph[a].append(b)
    graph[b].append(a)

# 次数1の頂点を根とする
deg = [len(x) for x in graph]
for i,x in enumerate(deg):
    if x == 1:
        break
root = i

parent = [0] * (N+1)
ts = [root]
stack = [root]
while stack:
    x = stack.pop()
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x
        stack.append(y)
        ts.append(y)

perm = np.zeros(N+K+10,np.int64)
perm[:K-1] = fact[K-2] * fact_inv[K-2::-1] % MOD

dp = [1] * (N+1)
for v in ts[::-1]:
    if v == root:
        break
    d = deg[v]
    dp[v] *= perm[d-1]; dp[v] %= MOD
    p = parent[v]
    dp[p] *= dp[v]; dp[p] %= MOD

x = dp[root]
x *= K*(K-1)
x %= MOD
print(x)