Problem p02985 (python)

Fixed submission:
import sys

sys.setrecursionlimit(10 ** 6)
mod = 10**9 + 7


def solve(now, from_, K, ab):
    if from_ == -1:
        can_num = K - 1
    else:
        can_num = K - 2
    if K < len(ab[now]):
        return 0
    case_num = 1
    for i in ab[now]:
        if i != from_:
            case_num *= can_num
            can_num -= 1
            case_num %= mod
    for i in ab[now]:
        if i != from_:
            case_num *= solve(i, now, K, ab)
            case_num %= mod
    return case_num


def main():
    N, K = map(int, input().split())
    if N == 1:
        print(K)
        exit(0)
    ab = [[] for _ in range(N)]
    for _ in range(N - 1):
        a, b = map(int, input().split())
        ab[a - 1].append(b - 1)
        ab[b - 1].append(a - 1)
    ans = K * solve(0, -1, K, ab)
    print(ans % mod)


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
mod = 10**9 + 7


def solve(now, from_, K, ab):
    if from_ == -1:
        can_num = K - 1
    else:
        can_num = K - 2
    if K < len(ab[now]):
        return 0
    case_num = 1
    for i in ab[now]:
        if i != from_:
            case_num *= can_num
            can_num -= 1
            case_num %= mod
    for i in ab[now]:
        if i != from_:
            case_num *= solve(i, now, K, ab)
            case_num %= mod
    return case_num


def main():
    N, K = map(int, input().split())
    if N == 1:
        print(K)
        exit(0)
    ab = [[] for _ in range(N)]
    for _ in range(N - 1):
        a, b = map(int, input().split())
        ab[a - 1].append(b - 1)
        ab[b - 1].append(a - 1)
    ans = K * solve(0, -1, K, ab)
    print(ans % mod)


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import numpy as np
mod = 10**9 + 7


def solve(now, from_, K, pa, b):
    if from_ == -1:
        can_num = K - 1
    else:
        can_num = K - 2
    if K < pa[now + 1] - pa[now]:
        return 0
    case_num = 1
    for i in range(pa[now], pa[now + 1]):
        if b[i] != from_:
            case_num *= can_num
            can_num -= 1
            case_num %= mod
    for i in range(pa[now], pa[now + 1]):
        if b[i] != from_:
            case_num *= solve(b[i], now, K, pa, b)
            case_num %= mod
    return case_num


def main():
    N, K = map(int, input().split())
    if N == 1:
        print(K)
        exit(0)
    ab = np.array([list(map(int, input().split())) for _ in range(N - 1)], dtype=np.int32)
    ab = (ab - 1).T
    ab = np.hstack((ab, ab[[1, 0]]))
    ab = ab[:, np.argsort(ab[0])]
    pa = np.unique(ab[0], return_index=True)[1]
    pa = np.append(pa, 2 * (N - 1))
    ans = K * solve(0, -1, K, pa, ab[1])
    print(ans % mod)


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import numpy as np
mod = 10**9 + 7


def solve(now, from_, K, pa, b):
    if from_ == -1:
        can_num = K - 1
    else:
        can_num = K - 2
    if K < pa[now + 1] - pa[now]:
        return 0
    case_num = 1
    for i in range(pa[now], pa[now + 1]):
        if b[i] != from_:
            case_num *= can_num
            can_num -= 1
            case_num %= mod
    for i in range(pa[now], pa[now + 1]):
        if b[i] != from_:
            case_num *= solve(b[i], now, K, pa, b)
            case_num %= mod
    return case_num


def main():
    N, K = map(int, input().split())
    ab = np.array([list(map(int, input().split())) for _ in range(N - 1)], dtype=np.int32)
    ab = (ab - 1).T
    ab = np.hstack((ab, ab[[1, 0]]))
    ab = ab[:, np.argsort(ab[0])]
    pa = np.unique(ab[0], return_index=True)[1]
    pa = np.append(pa, 2 * (N - 1))
    ans = K * solve(0, -1, K, pa, ab[1])
    print(ans % mod)


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import numpy as np
mod = 10**9 + 7


def solve(now, from_, K, pa, b):
    if from_ == -1:
        can_num = K - 1
    else:
        can_num = K - 2
    if K < pa[now + 1] - pa[now]:
        return 0
    case_num = 1
    for i in range(pa[now], pa[now + 1]):
        if b[i] != from_:
            case_num *= can_num
            can_num -= 1
            case_num %= mod
    for i in range(pa[now], pa[now + 1]):
        if b[i] != from_:
            case_num *= solve(b[i], now, K, pa, b)
            case_num %= mod
    return case_num


def main():
    N, K = map(int, input().split())
    ab = np.array([list(map(int, input().split())) for _ in range(N - 1)], dtype=np.int32)
    ab = ab - 1
    ab = np.hstack((ab, ab[[1, 0]]))
    ab = ab[:, np.argsort(ab[0])]
    pa = np.unique(ab[0], return_index=True)[1]
    pa = np.append(pa, 2 * (N - 1))
    ans = K * solve(0, -1, K, pa, ab[1])
    print(ans % mod)


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import sys
import numpy as np

i8 = np.int64
i4 = np.int32
mod = 10**9 + 7


def solve(now, from_, K, pa, b):
    if from_ == -1:
        can_num = K - 1
    else:
        can_num = K - 2
    if K < pa[now + 1] - pa[now]:
        return 0
    case_num = i8(1)
    for i in range(pa[now], pa[now + 1]):
        if b[i] != from_:
            case_num *= can_num
            can_num -= 1
            case_num %= mod
    for i in range(pa[now], pa[now + 1]):
        if b[i] != from_:
            case_num *= solve(b[i], now, K, pa, b)
            case_num %= mod
    return case_num


def main():
    pin = np.fromstring(sys.stdin.buffer.read(), dtype=i4, sep=' ')
    N = pin[0]
    K = pin[1]
    ab = pin[2:2 * N].reshape(N - 1, 2).T
    ab = ab - 1
    ab = np.hstack((ab, ab[[1, 0]]))
    ab = ab[:, np.argsort(ab[0])]
    pa = np.unique(ab[0], return_index=True)[1]
    pa = np.append(pa, 2 * (N - 1))
    ans = K * solve(0, -1, K, pa, ab[1])
    print(ans % mod)


if __name__ == "__main__":
    main()