Problem p02985 (python)

Fixed submission:
#頂点を見た時に，枝分かれの本数に応じてansを更新
#permutationのリストを先に作っておく

import sys
input = sys.stdin.readline

mod=10**9+7


N,K=map(int,input().split())
ab=[list(map(int, input().split())) for _ in range(N - 1)]
adj=[[] for _ in range(N)]

for a,b in ab:
    adj[a-1].append(b-1)
    adj[b-1].append(a-1)
    
    
#順列(K-2)Piを先に
P=[1]
for i in range(N):
    P.append((P[-1]*(K-2-i))%mod)
    

q=[(0,-1,0)]#node,parent,depth

ans=1
while q!=[]:
    v,p,d=q.pop()
    #深さ0or1で塗ることができる色数は異なる
    l=len(adj[v])
    if d==0:
        ans*=(K-1)
    ans*=(P[l-1])
    ans%=mod
    for i in range(l):
        next_v=adj[v][i]
        #葉は追加しない
        if next_v!=p and len(adj[next_v])!=1:
            q.append((next_v,v,d+1))
if N==1:
    ans=1
print((ans*K)%mod)

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

def I(): return int(input())
def MI(): return map(int, input().split())
def LI(): return list(map(int, input().split()))

def main():
    import queue
    mod=10**9+7
    N,K=MI()

    #0~Nまで逆元などを事前計算

    def calc(n, r, mod):
        if (r < 0) or (n < r):
            return 0
        return (fact[n] * factinv[n-r])%mod

    fact=[1,1]
    factinv=[1,1]
    inv=[0,1]

    for i in range(2, 10**5+5 + 1):
        fact.append((fact[-1] * i) % mod)
        inv.append((-inv[mod % i] * (mod // i)) % mod)
        factinv.append((factinv[-1] * inv[-1]) % mod)
    
    adj=[[]for _ in range(N)]
    for i in range(N-1):
        a,b=MI()
        a-=1
        b-=1
        adj[a].append(b)
        adj[b].append(a)
    
    #子要素を全部いっぺんに塗る
    ans=K#根の塗り方
    
    q=queue.Queue()
    q.put((0,0,-1))
    
    while not q.empty():
        v,d,p=q.get()
        cnt=0#子の個数
        for nv in adj[v]:
            if nv!=p:
                cnt+=1
                q.put((nv,d+1,v))
            
        if cnt!=0:#葉以外の場合は計算
            np=2#親方向で気にすべき個数
            if d==0:
                np=1#根の親はいない
                
        ans=(ans*calc(K-np,cnt,mod))%mod
        
    print(ans)

        
    

main()

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

def I(): return int(input())
def MI(): return map(int, input().split())
def LI(): return list(map(int, input().split()))

def main():
    import queue
    mod=10**9+7
    N,K=MI()

    #0~Nまで逆元などを事前計算

    def calc(n, r, mod):
        if (r < 0) or (n < r):
            return 0
        return (fact[n] * factinv[n-r])%mod

    fact=[1,1]
    factinv=[1,1]
    inv=[0,1]

    for i in range(2, N+5 + 1):
        fact.append((fact[-1] * i) % mod)
        inv.append((-inv[mod % i] * (mod // i)) % mod)
        factinv.append((factinv[-1] * inv[-1]) % mod)
    
    adj=[[]for _ in range(N)]
    for i in range(N-1):
        a,b=MI()
        a-=1
        b-=1
        adj[a].append(b)
        adj[b].append(a)
    
    #子要素を全部いっぺんに塗る
    ans=K#根の塗り方
    
    q=queue.Queue()
    q.put((0,0,-1))
    
    while not q.empty():
        v,d,p=q.get()
        cnt=0#子の個数
        for nv in adj[v]:
            if nv!=p:
                cnt+=1
                q.put((nv,d+1,v))
            
        if cnt!=0:#葉以外の場合は計算
            np=2#親方向で気にすべき個数
            if d==0:
                np=1#根の親はいない
                
        ans=(ans*calc(K-np,cnt,mod))%mod
        
    print(ans)

        
    

main()

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

def I(): return int(input())
def MI(): return map(int, input().split())
def LI(): return list(map(int, input().split()))

def main():
    import queue
    mod=10**9+7
    N,K=MI()

    #0~Nまで逆元などを事前計算

    def calc(n, r, mod):
        if (r < 0) or (n < r):
            return 0
        return (fact[n] * factinv[n-r])%mod

    fact=[1,1]
    factinv=[1,1]
    inv=[0,1]

    for i in range(2, K + 1):
        fact.append((fact[-1] * i) % mod)
        inv.append((-inv[mod % i] * (mod // i)) % mod)
        factinv.append((factinv[-1] * inv[-1]) % mod)
    
    adj=[[]for _ in range(N)]
    for i in range(N-1):
        a,b=MI()
        a-=1
        b-=1
        adj[a].append(b)
        adj[b].append(a)
    
    #子要素を全部いっぺんに塗る
    ans=K#根の塗り方
    
    q=queue.Queue()
    q.put((0,0,-1))
    
    while not q.empty():
        v,d,p=q.get()
        cnt=0#子の個数
        for nv in adj[v]:
            if nv!=p:
                cnt+=1
                q.put((nv,d+1,v))
            
        if cnt!=0:#葉以外の場合は計算
            np=2#親方向で気にすべき個数
            if d==0:
                np=1#根の親はいない
                
        ans=(ans*calc(K-np,cnt,mod))%mod
        
    print(ans)

        
    

main()