Problem p02685 (python)

Fixed submission:
import numpy as np


class mod_calc():
    def __init__(self, divisor, max_n):
        # assume prime
        self._divisor = divisor
        # nsq = math.ceil(math.sqrt(max_n))
        # feature fix_n? max_k?
        fact = np.empty((max_n+1,), dtype=int)
        fact[1] = fact[0] = f = 1
        for i in range(2, len(fact)):
            f = (f * i) % divisor
            fact[i] = f
        self._fact = fact
        fact_inv = np.empty((max_n+1,), dtype=int)
        fact_inv[max_n] = fi = pow(int(fact[max_n]), divisor-2, divisor)
        for i in range(max_n, 2, -1):
            fi = (fi * i) % divisor
            fact_inv[i-1] = fi
        fact_inv[1] = fact_inv[0] = 1
        self._fact_inv = fact_inv

    def comb(self, n, k):
        k = min(k, n-k)
        if k == 0 or k == n:
            return 1
        return self.mul(self._fact[n], self._fact_inv[n-k], self._fact_inv[k])

    def pow(self, a, n):
        return pow(a, n, self._divisor)

    def mul(self, *a):
        # TODO montogomery
        # https://ikatakos.com/pot/programming_algorithm/number_theory/montgomery_multiplication
        r = a[0]
        for m in a[1:]:
            r = (r * m) % self._divisor
        return r

    def sum(self, *a):
        r = a[0]
        for m in a[1:]:
            r = (r + m) % self._divisor
        return r


def main():
    N, M, K = [int(x) for x in input().split()]
    if M == 1:
        if K == N - 1:
            return 1
        else:
            return 0
    if N == 1:
        return M
    c = mod_calc(998244353, N-1)
    count = 0
    for k in range(K+1):
        count = c.sum(c.mul(c.comb(N-1,k), M, c.pow(M-1, N-k-1)), count)
    return count


if __name__ == "__main__":
    print(main())

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


class mod_calc():
    def __init__(self, divisor, max_n):
        # assume prime
        self._divisor = divisor
        # nsq = math.ceil(math.sqrt(max_n))
        # feature fix_n? max_k?
        fact = np.empty((max_n+1,), dtype=int)
        fact[1] = fact[0] = f = 1
        for i in range(2, len(fact)):
            f = (f * i) % divisor
            fact[i] = f
        self._fact = fact
        fact_inv = np.empty((max_n+1,), dtype=int)
        fact_inv[max_n] = fi = pow(int(fact[max_n]), divisor-2, divisor)
        for i in range(max_n, 2, -1):
            fi = (fi * i) % divisor
            fact_inv[i-1] = fi
        fact_inv[1] = fact_inv[0] = 1
        self._fact_inv = fact_inv

    def comb(self, n, k):
        k = min(k, n-k)
        if k == 0 or k == n:
            return 1
        return self.mul(self._fact[n], self._fact_inv[n-k], self._fact_inv[k])

    def pow(self, a, n):
        return pow(a, n, self._divisor)

    def mul(self, *a):
        # TODO montogomery
        # https://ikatakos.com/pot/programming_algorithm/number_theory/montgomery_multiplication
        r = a[0]
        for m in a[1:]:
            r = (r * m) % self._divisor
        return r

    def sum(self, *a):
        r = a[0]
        for m in a[1:]:
            r = (r + m) % self._divisor
        return r


def main():
    N, M, K = [int(x) for x in input().split()]
    if M == 1:
        return 1
    c = mod_calc(998244353, N-1)
    count = 0
    for k in range(K+1):
        count = c.sum(c.mul(c.comb(N-1,k), M, c.pow(M-1, N-k-1)), count)
    return count


if __name__ == "__main__":
    print(main())