Problem p02685 (python)

Fixed submission:
N,M,K = map(int,input().split())

p = 998244353
# 何個を同じ色の組にするか0~K
cnt = 0
com = 1
for k in range(K+1):
    l = N-k
    # N-1個の組のうちどれを同じ色にするかN-1Ck
    # それぞれで、N-k個を別の色の順列にする
    # M*(M-1)**(N-k-1)
    x = pow((M-1),(l-1),p)
    cnt += M * x * com
    cnt %= p

    # 次のkのnCk
    com *= N-(k+1)
    com *= pow(k+1,p-2,p)
    com %= p


print(cnt)

Buggy submission (verdict: Runtime Error):
class Combination:
    """
    O(n)の前計算を1回行うことで，O(1)でnCr mod mを求められる
    n_max = 10**6のとき前処理は約950ms (PyPyなら約340ms, 10**7で約1800ms)
    使用例：
    comb = Combination(1000000)
    print(comb(5, 3))  # 10
    """
    def __init__(self, n_max, mod=10**9+7):
        self.mod = mod
        self.modinv = self.make_modinv_list(n_max)
        self.fac, self.facinv = self.make_factorial_list(n_max)

    def __call__(self, n, r):
        return self.fac[n] * self.facinv[r] % self.mod * self.facinv[n-r] % self.mod

    def make_factorial_list(self, n):
        # 階乗のリストと階乗のmod逆元のリストを返す O(n)
        # self.make_modinv_list()が先に実行されている必要がある
        fac = [1]
        facinv = [1]
        for i in range(1, n+1):
            fac.append(fac[i-1] * i % self.mod)
            facinv.append(facinv[i-1] * self.modinv[i] % self.mod)
        return fac, facinv

    def make_modinv_list(self, n):
        # 0からnまでのmod逆元のリストを返す O(n)
        modinv = [0] * (n+1)
        modinv[1] = 1
        for i in range(2, n+1):
            modinv[i] = self.mod - self.mod//i * modinv[self.mod%i] % self.mod
        return modinv

N,M,K = map(int,input().split())

p = 998244353

comb = Combination(N-1,p)

# 何個を同じ色の組にするか0~K
cnt = 0
for k in range(K+1):
    buf = 0
    # N-1個の組のうちどれを同じ色にするかN-1Ck
    buf = comb(N-1,k)

    # それぞれで、N-k個を別の色の順列にする
    # M*(M-1)**(N-k-1)
    buf *= M%p
    buf %= p
    buf *= pow((M-1),(N-k-1),p)
    cnt += buf%p
    cnt %= p

print(cnt)

Buggy submission (verdict: Runtime Error):
class Combination:
    """
    O(n)の前計算を1回行うことで，O(1)でnCr mod mを求められる
    n_max = 10**6のとき前処理は約950ms (PyPyなら約340ms, 10**7で約1800ms)
    使用例：
    comb = Combination(1000000)
    print(comb(5, 3))  # 10
    """
    def __init__(self, n_max, mod=10**9+7):
        self.mod = mod
        self.modinv = self.make_modinv_list(n_max)
        self.fac, self.facinv = self.make_factorial_list(n_max)

    def __call__(self, n, r):
        return self.fac[n] * self.facinv[r] % self.mod * self.facinv[n-r] % self.mod

    def make_factorial_list(self, n):
        # 階乗のリストと階乗のmod逆元のリストを返す O(n)
        # self.make_modinv_list()が先に実行されている必要がある
        fac = [1]
        facinv = [1]
        for i in range(1, n+1):
            fac.append(fac[i-1] * i % self.mod)
            facinv.append(facinv[i-1] * self.modinv[i] % self.mod)
        return fac, facinv

    def make_modinv_list(self, n):
        # 0からnまでのmod逆元のリストを返す O(n)
        modinv = [0] * (n+1)
        modinv[1] = 1
        for i in range(2, n+1):
            modinv[i] = self.mod - self.mod//i * modinv[self.mod%i] % self.mod
        return modinv

N,M,K = map(int,input().split())

p = 998244353

comb = Combination(N-1,p)

# 何個を同じ色の組にするか0~K
cnt = 0
for k in range(K+1):
    buf = 0
    # N-1個の組のうちどれを同じ色にするかN-1Ck
    buf = comb(N-1,k)

    # それぞれで、N-k個を別の色の順列にする
    # M*(M-1)**(N-k-1)
    if N-k-1 >= 0:
        buf *= M
        buf *= pow((M-1),(N-k-1),p)
        buf %= p
        cnt += buf
        cnt %= p

print(cnt)

Buggy submission (verdict: Runtime Error):
class Combination:
    """
    O(n)の前計算を1回行うことで，O(1)でnCr mod mを求められる
    n_max = 10**6のとき前処理は約950ms (PyPyなら約340ms, 10**7で約1800ms)
    使用例：
    comb = Combination(1000000)
    print(comb(5, 3))  # 10
    """
    def __init__(self, n_max, mod=10**9+7):
        self.mod = mod
        self.modinv = self.make_modinv_list(n_max)
        self.fac, self.facinv = self.make_factorial_list(n_max)

    def __call__(self, n, r):
        return self.fac[n] * self.facinv[r] % self.mod * self.facinv[n-r] % self.mod

    def make_factorial_list(self, n):
        # 階乗のリストと階乗のmod逆元のリストを返す O(n)
        # self.make_modinv_list()が先に実行されている必要がある
        fac = [1]
        facinv = [1]
        for i in range(1, n+1):
            fac.append(fac[i-1] * i % self.mod)
            facinv.append(facinv[i-1] * self.modinv[i] % self.mod)
        return fac, facinv

    def make_modinv_list(self, n):
        # 0からnまでのmod逆元のリストを返す O(n)
        modinv = [0] * (n+1)
        modinv[1] = 1
        for i in range(2, n+1):
            modinv[i] = self.mod - self.mod//i * modinv[self.mod%i] % self.mod
        return modinv

N,M,K = map(int,input().split())

p = 998244353

comb = Combination(N-1,p)

# 何個を同じ色の組にするか0~K
cnt = 0
for k in range(K+1):
    buf = 0
    # N-1個の組のうちどれを同じ色にするかN-1Ck
    buf = comb(N-1,k)

    # それぞれで、N-k個を別の色の順列にする
    # M*(M-1)**(N-k-1)
    buf *= M
    buf *= pow((M-1),(N-k-1),p)
    buf %= p
    cnt += buf
    cnt %= p

print(cnt)