Problem p02985 (python)

Fixed submission:
from collections import deque


def solve(N, K, ABs):
    if N == 1:
        return K

    Ds = [-1 for _ in range(N + 1)]
    Ws = {i: [] for i in range(N + 1)}
    for a, b in ABs:
        Ws[a].append(b)
        Ws[b].append(a)

    for i in range(1, N + 1):
        if Ds[i] == -1:
            Ds[i] = K

            if Ws[i]:
                s = Ws[i][0]
                Ds[s] = K - 1

                queue = deque()
                queue.append(i)
                queue.append(s)

                while queue:
                    s = queue.popleft()
                    k = K - 2
                    for w in Ws[s]:
                        if Ds[w] == -1:
                            Ds[w] = k
                            k -= 1
                            queue.append(w)

    ans = 1
    for i in range(1, N + 1):
        if Ds[i] <= 0:
            return 0
        ans *= Ds[i]
        ans %= 1000000007
    return ans % 1000000007

if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    ABs = [tuple(map(int, input().split(" "))) for _ in range(N - 1)]
    print(solve(N, K, ABs))

Buggy submission (verdict: Runtime Error):
from collections import deque
 
 
def solve(N, K, ABs):
     if N == 1:
        return K

    Ds = [0 for _ in range(N + 1)]
    Ws = {i: [] for i in range(N + 1)}
    for a, b in ABs:
        Ws[a].append(b)
        Ws[b].append(a)
 
    Ds[1] = K
 
    s = Ws[1][0]
    Ds[s] = K - 1
 
    queue = deque()
    queue.append(1)
    queue.append(s)
 
    while queue:
        s = queue.popleft()
        k = K - 2
        for w in Ws[s]:
            if Ds[w] == 0:
                Ds[w] = k
                k -= 1
                queue.append(w)
    ans = 1
    for i in range(1, N + 1):
        if Ds[i] <= 0:
            return 0
        ans *= Ds[i]
        ans %= 1000000007
    return (ans % 1000000007)
 
if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    ABs = [tuple(map(int, input().split(" "))) for _ in range(N - 1)]
    print(solve(N, K, ABs))

Buggy submission (verdict: Runtime Error):
from collections import deque


def solve(N, K, ABs):

    Ds = [0 for _ in range(N + 1)]
    Ws = {i: [] for i in range(N + 1)}
    for a, b in ABs:
        Ws[a].append(b)
        Ws[b].append(a)

    Ds[1] = K

    s = Ws[1][0]
    Ds[s] = K - 1

    queue = deque()
    queue.append(1)
    queue.append(s)

    while queue:
        s = queue.popleft()
        k = K - 2
        for w in Ws[s]:
            if Ds[w] == 0:
                Ds[w] = k
                k -= 1
                queue.append(w)
    ans = 1
    for i in range(1, N + 1):
        if Ds[i] <= 0:
            return 0
        ans *= Ds[i]
        ans %= 1000000007
    return (ans % 1000000007)

if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    ABs = [tuple(map(int, input().split(" "))) for _ in range(N - 1)]
    print(solve(N, K, ABs))

Buggy submission (verdict: Runtime Error):
from collections import deque


def solve(N, K, ABs):
    Ds = [0 for _ in range(N + 1)]
    Ws = {i: [] for i in range(N + 1)}
    for a, b in ABs:
        Ws[a].append(b)
        Ws[b].append(a)

    Ds[1] = K

    s = Ws[1][0]
    Ds[s] = K - 1

    queue = deque()
    queue.append(1)
    queue.append(s)

    while queue:
        s = queue.popleft()
        k = K - 2
        for w in Ws[s]:
            if Ds[w] == 0:
                Ds[w] = k
                k -= 1
                queue.append(w)
    ans = 1
    for i in range(1, N + 1):
        if Ds[i] <= 0:
            return 0
        ans *= Ds[i]
        ans %= 1000000007
    return (ans % 1000000007)


# print(solve(4, 3, [(1, 2), (2, 3), (3, 4)]))
assert (solve(1, 1, []) == 1)
assert (solve(1, 2, []) == 2)
assert (solve(1, 10000, []) == 10000)

assert (solve(2, 1, [(1, 2)]) == 0)
assert (solve(3, 1, [(1, 2), (1, 3)]) == 0)
assert (solve(3, 2, [(1, 2), (1, 3)]) == 0)
assert (solve(3, 2, [(1, 2), (2, 3)]) == 0)
assert (solve(2, 2, [(1, 2)]) == 2)

assert (solve(2, 10000, [(1, 2)]) == 10000 * 9999)
assert (solve(3, 10000, [(1, 2), (1, 3)]) == (10000 * 9999 * 9998) % 1000000007)
assert (solve(3, 10000, [(1, 2), (2, 3)]) == (10000 * 9999 * 9998) % 1000000007)

assert (solve(4, 3, [(1, 2), (1, 3), (1, 4)]) == 0)

assert (solve(4, 3, [(1, 2), (2, 3), (3, 4)]) == 6)
assert (solve(5, 4, [(1, 2), (1, 3), (1, 4), (4, 5)]) == 48)
assert (solve(16, 22,
              [(12, 1), (3, 1), (4, 16), (7, 12), (6, 2), (2, 15), (5, 16), (14, 16), (10, 11), (3, 10), (3, 13),
               (8, 6), (16, 8), (9, 12), (4, 3)]) == 271414432)

if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    ABs = [tuple(map(int, input().split(" "))) for _ in range(N - 1)]
    print(solve(N, K, ABs))

Buggy submission (verdict: Runtime Error):
from collections import deque


def solve(N, K, ABs):
    # print(N, ",", K, ", ", ABs)
    Ds = [0 for _ in range(N + 1)]
    Ws = {i: [] for i in range(N + 1)}
    for a, b in ABs:
        Ws[a].append(b)
        Ws[b].append(a)

    Ds[1] = K
    s = Ws[1][0]
    Ds[s] = K - 1

    queue = deque()
    queue.append(1)
    queue.append(s)

    while queue:
        s = queue.popleft()
        k = K - 2
        for w in Ws[s]:
            if Ds[w] == 0:
                Ds[w] = k
                k -= 1
                queue.append(w)
    ans = 1
    for i in range(1, N + 1):
        ans *= Ds[i]
        ans %= 1000000007
    return ans


# print(solve(4, 3, [(1, 2), (2, 3), (3, 4)]))
# assert (solve(4, 3, [(1, 2), (2, 3), (3, 4)]) == 6)
# assert(solve(5, 4, [(1, 2), (1, 3), (1, 4), (4, 5)]) == 48)

if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    ABs = [tuple(map(int, input().split(" "))) for _ in range(N - 1)]
    print(solve(N, K, ABs))