Problem p02985 (python)

Fixed submission:
import sys
input = sys.stdin.readline
from collections import *

MOD = 10**9+7

class Comb:
    def __init__(self, N):
        self.fac = [1] * (N+1)
        for i in range(2, N+1): self.fac[i] = self.fac[i-1] * i % MOD

    def pow(self, a, b):
        res = 1
        while b:
            if b & 1: res = res * a % MOD
            a = a**2 % MOD
            b >>= 1
        return res

    def comb(self, n, r):
        if r < 0 or r > n: return 0
        return (self.fac[n] * self.pow(self.fac[r], MOD-2)) % MOD * self.pow(self.fac[n-r], MOD-2) % MOD

    def permutation(self, n, r):
        return self.fac[n] * self.pow(self.fac[n-r], MOD-2) % MOD

def main():
    N, K = list(map(int, input().split()))
    g = [[] for i in range(N+1)]
    ans = K

    for i in range(N-1):
        a, b = list(map(int, input().split()))
        g[a].append(b)
        g[b].append(a)

    comb = Comb(K)

    q = deque([(-1, 1)])
    while q:
        par, chi = q.pop()
        r = len(g[chi])-1
        if par == -1: k = K-1; r += 1
        else: k = K-2
        if k < r: ans = 0; break

        ans = ans * comb.permutation(k, r) % MOD
        for c in g[chi]:
            if c != par: q.append((chi, c))

    print(ans%MOD)
    
if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline

MOD = 10**9+7

class Comb:
    def __init__(self, N):
        self.fac = [1] * (N+1)
        for i in range(2, N+1): self.fac[i] = self.fac[i-1] * i % MOD

    def pow(self, a, b):
        res = 1
        while b:
            if b & 1: res = res * a % MOD
            a = a**2 % MOD
            b >>= 1
        return res

    def comb(self, n, r):
        if r < 0 or r > n: return 0
        return (self.fac[n] * self.pow(self.fac[r], MOD-2)) % MOD * self.pow(self.fac[n-r], MOD-2) % MOD

    def permutation(self, n, r):
        return self.fac[n] * self.pow(self.fac[n-r], MOD-2) % MOD

def dfs(par, chi, d):
    global ans, comb
    
    r = len(g[chi])-1
    if d == 0: k = K-1; r += 1
    else: k = K-2

    ans = ans * comb.permutation(k, r) % MOD
    for c in g[chi]:
        if c == par: continue
        dfs(chi, c, d+1)


N, K = list(map(int, input().split()))
g = [[] for i in range(N+1)]
ans = K

for i in range(N-1):
    a, b = list(map(int, input().split()))
    g[a].append(b)
    g[b].append(a)

comb = Comb(K)

dfs(-1, 1, 0)

print(ans%MOD)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline

MOD = 10**9+7

def pow(a):
    b = MOD-2
    res = 1
    while b:
        if b & 1: res = res * a % MOD
        a = a**2 % MOD
        b >>= 1
    return res

def permutation(n, r):
    if r == 0: return 1
    return fac[n] * pow(fac[n-r]) % MOD

def dfs(v, p):
    if p == -1: n = permutation(K-1, len(g[v]))
    else: n = permutation(K-2, len(g[v])-1)
    
    for u in g[v]:
        if u == p: continue
        n = n * dfs(u, v) % MOD
    return n

N, K = map(int, input().split())
g = [[] for i in range(N)]

for i in range(N-1):
    a, b = map(int, input().split())
    a -= 1; b -= 1
    g[a].append(b)
    g[b].append(a)
    
fac = [1] * (K+1)

for i in range(2, K+1): fac[i] = fac[i-1] * i % MOD

n = dfs(0, -1)

print(K*n % MOD)