ABC 167 E - Colorful Blocks
問題
一列に並んだ $n$ 個のブロックを、同じ色で塗られた隣り合うブロックの組が $k$ 組以下になるように $ m $ 色で塗る塗り方は何通りあるか。（$n,m \leq 2 \times 10 ^ 5, \ k \leq n - 1$）
解説
まず、同じ色で塗られた隣り合うブロックの組が $i$ 組になる塗り方について考える。隣り合うブロックの組 $n - 1$ 個から同じ色で塗る組 $i$ 個を選ぶ組合せは $ _ {n-1} C _ {i}$ 通りある。同じ色で塗る隣り合うブロックの組を決めるとブロック全体は $n-i$ 個のグループに分かれ、各グループに属するブロックは同じ色で塗り、隣り合うグループでは異なる色で塗ることになる。$n-i$ 個のグループを隣り合うグループが同じ色にならないように $ m $ 色で塗り分ける塗り方は、最初のグループを塗る色が $ m $ 通りあり、それ以降のグループは直前に使った色以外から選ぶので $ m - 1 $ 通りで、$ m \times (m - 1) ^ {n - i -1} $ 通りとなる。よって、同じ色で塗られた隣り合うブロックの組が $i$ 組になる塗り方は $_ {n-1} C _ {i} \times m \times (m - 1) ^ {n - i -1}$ 通りで、答えはこれを $i = 0$ から $k$ まで合計するので、 $$ \sum _ {i = 0} ^ {k} \ _ {n-1} C _ {i} \times m \times (m - 1) ^ {n - i -1} $$
実装
二項係数を繰り返し計算するため、あらかじめ $\mathrm{fct}[i] = i !,\mathrm{inv}[i] = (i !) ^ {-1} \ (\mathrm{mod}\ p)$ となる配列 $\mathrm{fct}$, $\mathrm{inv}$ を用意し、 $$ _ {n-1} C _ {i} = \mathrm{fct}[n - 1] \times \mathrm{inv}[i] \times \mathrm{inv}[n - 1 - i] $$ として計算した。
n, m, k = map(int, input().split()) p = 998244353 fct, inv = [1], [1] for i in range(1, n): fct.append((fct[-1] * i) % p) inv.append((inv[-1] * pow(i, p - 2, p)) % p) ans = 0 for i in range(k + 1): cmb = fct[n - 1] * inv[i] * inv[n - 1 - i] ans = (ans + cmb * m * pow(m - 1, n - 1 - i, p)) % p print(ans)
