Problem: There are N blocks arranged in a row. Let us paint these blocks.  We will consider two ways to paint the blocks different if and only if there is a block painted in different colors in those two ways.  Find the number of ways to paint the blocks under the following conditions:    * For each block, use one of the M colors, Color 1 through Color M, to paint it. It is not mandatory to use all the colors.   * There may be at most K pairs of adjacent blocks that are painted in the same color.  Since the count may be enormous, print it modulo 998244353.

Submission attempts (in order):
import sys
import itertools
import numpy as np

MOD = 998244353
MAX = 10 ** 5 * 2 + 5

fac = [0 for i in range(MAX)]
finv = [0 for i in range(MAX)]
inv = [0 for i in range(MAX)]

def comInit(mod):
    fac[0], fac[1] = 1, 1
    finv[0], finv[1] = 1, 1
    inv[1] = 1
    for i in range(2, MAX):
        fac[i] = fac[i - 1] * i % mod
        inv[i] = mod - inv[mod % i] * (mod // i) % mod
        finv[i] = finv[i - 1] * inv[i] % mod

def com(n, r, mod):
    if n < r:
        return 0
    if n < 0 or r < 0:
        return 0
    return fac[n] * (finv[r] * finv[n - r] % mod) % mod

comInit(MOD)

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
 
n, m, k = map(int, readline().split())

ans = 0
for x in range(k, -1, -1):
    ans += now * com(n - 1, x, MOD) % MOD
    # ans += tmp
print(ans % MOD)

import sys
import itertools
import numpy as np

MOD = 998244353
MAX = 10 ** 5 * 2 + 5

fac = [0 for i in range(MAX)]
finv = [0 for i in range(MAX)]
inv = [0 for i in range(MAX)]

def comInit(mod):
    fac[0], fac[1] = 1, 1
    finv[0], finv[1] = 1, 1
    inv[1] = 1
    for i in range(2, MAX):
        fac[i] = fac[i - 1] * i % mod
        inv[i] = mod - inv[mod % i] * (mod // i) % mod
        finv[i] = finv[i - 1] * inv[i] % mod

def com(n, r, mod):
    if n < r:
        return 0
    if n < 0 or r < 0:
        return 0
    return fac[n] * (finv[r] * finv[n - r] % mod) % mod

comInit(MOD)

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
 
n, m, k = map(int, readline().split())

ans = 0
for x in range(k, -1, -1):
    # now *= (m - 1) % MOD
    ans += com(n - 1, x, MOD) % MOD
    # ans += tmp
print(ans % MOD)

import sys
import itertools
import numpy as np

MOD = 998244353
MAX = 10 ** 5 * 2 + 5

fac = [0 for i in range(MAX)]
finv = [0 for i in range(MAX)]
inv = [0 for i in range(MAX)]

def comInit(mod):
    fac[0], fac[1] = 1, 1
    finv[0], finv[1] = 1, 1
    inv[1] = 1
    for i in range(2, MAX):
        fac[i] = fac[i - 1] * i % mod
        inv[i] = mod - inv[mod % i] * (mod // i) % mod
        finv[i] = finv[i - 1] * inv[i] % mod

def com(n, r, mod):
    if n < r:
        return 0
    if n < 0 or r < 0:
        return 0
    return fac[n] * (finv[r] * finv[n - r] % mod) % mod

comInit(MOD)

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
 
n, m, k = map(int, readline().split())

ans = 0
col = m
for x in range(n - 1, -1, -1):
    now = col
    now *= com(n - 1, x, MOD) % MOD
    if x <= k:
        ans += now
    col = (col * (m - 1)) % MOD

print(ans % MOD)

Verdicts: Runtime Error

Wrong Answer

Accepted