Problem p02985 (python)

Fixed submission:


MOD = 10**9 + 7
N,K = map(int, input().split())
es = [[] for _ in range(N)]
for _ in range(N-1):
    a,b = map(int, input().split())
    es[a-1].append(b-1)
    es[b-1].append(a-1)

MAXN = K+5
fac = [1,1] + [0]*MAXN
finv = [1,1] + [0]*MAXN
inv = [0,1] + [0]*MAXN
for i in range(2,MAXN+2):
    fac[i] = fac[i-1] * i % MOD
    inv[i] = -inv[MOD%i] * (MOD // i) % MOD
    finv[i] = finv[i-1] * inv[i] % MOD

def nPr(n,r):
    if n < r:
        return 0
    return fac[n] * finv[n-r]

# 「距離が２以内は異なる色 = あるノードとそれに隣接するノードはそれぞれ異なる色」の言い換えがポイント
# 根はK通り、根の子の数をmとすると、根の子の塗り方は、K-1個の中からm個並べる方法なので、nPr(K-1, m)通り
# それ以降についてあるノードを見るとき、そのノードと親には既に色が塗られている場合を考える。
# 親とその親の２色は使われてるから、そのノードのこの塗り方はK-2のうちから子の数だけ並べるときなので、nPr(K-2, 子の数)


ans = (K * nPr(K-1, len(es[0]))) % MOD
stack = []
for nxt in es[0]:
    stack.append(nxt)

checked = [False] * N
checked[0] = True

# ノード0の子から見ていって、その子の塗り方をかけ合わせていく
while stack:
    curr = stack.pop()
    checked[curr] = True
    # currの親はすでに塗ってあるので、r = len(es[curr])-1
    ans *= nPr(K-2, len(es[curr])-1)
    ans %= MOD
    for nxt in es[curr]:
        if not checked[nxt]:
            stack.append(nxt)

print(ans % MOD)

Buggy submission (verdict: Runtime Error):
N,K = map(int, input().split())
es = [[] for _ in range(N)]
for _ in range(N-1):
    a,b = map(int, input().split())
    es[a-1].append(b-1)
    es[b-1].append(a-1)


MOD = 10**9+7
 
MAXN = N+5
fac = [1,1] + [0]*MAXN
finv = [1,1] + [0]*MAXN
inv = [0,1] + [0]*MAXN
for i in range(2,MAXN+2):
    fac[i] = fac[i-1] * i % MOD
    inv[i] = -inv[MOD%i] * (MOD // i) % MOD
    finv[i] = finv[i-1] * inv[i] % MOD
 

def nPr(n,r):
    if n < r:
        return 0
    return fac[n] * finv[n-r]


# 始点のノードがK個の色でぬれて、始点の子供がnPr(K-1, 始点の子の数)通りの塗り分け方がある。
# それらの子については、nPr(K-2, 子の数)の塗り方。

q  = []
used = [False] * N
used[0] = True
for nxt in es[0]:
    q.append(nxt)
    used[nxt] = True


ans = K * nPr(K-1, len(es[0]))
ans %= MOD
while q:
    curr = q.pop()
    ans *= nPr(K-2, len(es[curr]) - 1)
    ans %= MOD
    for nxt in es[curr]:
        if used[nxt]:
            continue

        used[nxt] = True
        q.append(nxt)


print(ans)

Buggy submission (verdict: Runtime Error):




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

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


MOD = 10**9+7
 
MAXN = N+5
fac = [1,1] + [0]*MAXN
finv = [1,1] + [0]*MAXN
inv = [0,1] + [0]*MAXN
for i in range(2,MAXN+2):
    fac[i] = fac[i-1] * i % MOD
    inv[i] = -inv[MOD%i] * (MOD // i) % MOD
    finv[i] = finv[i-1] * inv[i] % MOD
 
def nPr(n,r):
    if n < r:
        return 0
    return fac[n] * finv[n-r]

ans = 1

checked = [False] * N
q = []
# 探索開始ノードの取りうる色K個と、そこに繋がるノードの色の取り方
ans *= K * nPr(K-1, len(es[0]))

checked[0] = True
for nxt in es[0]:
    q.append(nxt)
    checked[nxt] = True

while q:
    curr = q.pop()
    ans *= nPr(K-2, len(es[curr]) - 1)
    ans %= MOD
    for nxt in es[curr]:
        if not checked[nxt]:
            q.append(nxt)
            checked[nxt] = True
print(ans % MOD)

Buggy submission (verdict: Runtime Error):




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

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


MOD = 10**9+7
 
MAXN = N+5
fac = [1,1] + [0]*MAXN
finv = [1,1] + [0]*MAXN
inv = [0,1] + [0]*MAXN
for i in range(2,MAXN+2):
    fac[i] = fac[i-1] * i % MOD
    inv[i] = -inv[MOD%i] * (MOD // i) % MOD
    finv[i] = finv[i-1] * inv[i] % MOD
 
def nPr(n,r):
    if n < r:
        return 0
    return fac[n] * finv[n-r]

ans = 1

checked = [False] * N
q = []
# 探索開始ノードの取りうる色K個と、そこに繋がるノードの色の取り方
ans *= K * nPr(K-1, len(es[0]))

checked[0] = True
for nxt in es[0]:
    q.append(nxt)
    checked[nxt] = True

while q:
    curr = q.pop()
    ans *= nPr(K-2, len(es[curr]) - 1)
    ans %= MOD
    for nxt in es[curr]:
        if not checked[nxt]:
            q.append(nxt)
            checked[nxt] = True
print(ans)

Buggy submission (verdict: Runtime Error):



N,K = map(int, input().split())
MOD = 10**9 + 7
es = [[] for _ in range(N)]
for _ in range(N-1):
    a,b = map(int, input().split())
    es[a-1].append(b-1)
    es[b-1].append(a-1)



checked = [False] * N

colors = [None] * N

def solve(curr, step):
    # step=0のとき、そのノードの色を何個選べるか確認する
    if step == 0:
        used = 0
        for nxt in es[curr]:
            used += solve(nxt, 1)
        colors[curr] = K - used
        # 次を探索する
        for nxt in es[curr]:
            if colors[nxt] == None:
                solve(nxt, 0)
    # step=1のとき、自分とその次のノードについて色が何個選べるか分かっているノード数を戻す
    elif step == 1:
        used = 0
        if colors[curr] != None:
            used += 1
        for nxt in es[curr]:
            used += solve(nxt, 2)
        return used

    # step=2の場合、自分について色が何個選べるか分かっている場合に１を戻す
    elif step == 2:
        if colors[curr] != None:
            return 1
        else:
            return 0

solve(0,0)

#print(colors)

ans = 1
for c in colors:
    ans *= c
    ans %= MOD

print(ans)