問題概要
正整数 $N, M, K$ が与えられる．
最初，$N$ 個のブロックが横一列に（隣接して）並んでいる．このブロックたちに色を塗りたい．次の条件を満たす色の塗り方を $\pmod{ 998244353 }$ で求めよ．
- それぞれのブロックは，色 $1$ から色 $M $ の $M $ 種類の色の内ちょうど一色で塗られる
- 隣接するブロックの組であって，同じ色で塗られているようなものの数は，高々 $K$ 個
なお，2 つのブロックの塗り方が異なるとは，あるブロックであって，片方の塗り方での色と他方での塗り方での色が異なるものが存在することを言う．
制約
- $2 \leq N, M \leq 2 \times 10^5$
- $0 \leq K \leq N - 1$
準備
本問題では，重複組合せを使って解を求めていきます．記法として，組合せ (Combination) を高校などでおなじみの $_n \mathrm C_k$ ではなく $\dbinom n k$ と表記します（主に書くのが楽になるからです）．同様に，重複組合せ $_n \mathrm{ H }_r$ は $\left( \dbinom n r \right)$ と書きます．なお，$\left( \dbinom n r \right) = \dbinom { n + r - 1 } r$ です*1．
考察
解法の前に，考察の一環として思い付きがちな誤解法について述べておきます（本番中，ちょっと考えて棄却する流れで進めたので）．$\mathrm{ mod }$ での数え上げといえば，非常によく出てくる解法は，DP か，四則演算ベースでえいやっとやる方法です．今回，DP の方はやや筋が悪いです．例えば，次のような DP
- dp[ （左から）何個のブロックに色を塗り終えたか ][ 色が隣接している箇所の数 ] := ways
は，$NK$ がすでに大きく，ややつらいです（遷移については割愛します）．ここで，難しくしているファクターは何かと考えると，同色が隣接する個数に関する条件です．この条件を一旦改変することで，何か知見が得られないか考えてみます（という流れで本番中考えてました，という話です．全体的に）．
解法
ということで，まずは $N$ 個のブロックを同色の隣接を完全に禁止して塗るバージョンについて考えてみます．この場合は，
- 先頭のブロックには $M $ 色から好きな色を選んで使える
- それ以外のブロックには，直前のブロックと被らなければよいので，$M - 1$ 色から好きな色を選べる
となります．被るかどうかだけに着目しているので，具体的にどの色なのかを考えなくてよいのはポイントかもしれません．で，この場合の答えの求め方ですが，
- 先頭は $M $ 通りなのでそのまま $M $
- 残りは，$N - 1$ 箇所に対して $M - 1$ 通りのチョイスを繰り返すので $( M - 1 )^{ N - 1 }$
となって，全体では $M ( M - 1 )^{ N - 1 }$ となります．
ここで，$N$ の値を色々変えてあげると，任意の $N'$ ($2 \leq N’ \leq N$) について同様の問題を解けることに気が付きます．それらの解たちを使って，本来解きたかった問題の答えを求めることができます．まずは $K$ が小さいケースについて考察してみましょう．
最初は $K = 1$ のケースです*2．この場合，同じ色のブロックが隣接する部分が一箇所ありますが，それを $1$ つのブロックに潰して（「縮約」的な感じ）しまった状態をイメージしてみると，$N - 1$ 個のブロックを隣接を許さず塗ったように見えます．逆に，$N - 1$ 個が塗られた状態から潰されたブロックを元に戻そうと考えると，戻せるブロックが $N - 1$ 箇所あり，そのどれを膨らませても違う塗り方が得られます．よって，$K = 1$ の場合の塗り方は，先程の式の $N$ を $N - 1$ に置き換えた上で $N - 1$ をかけて，$M ( M - 1 )^{ N - 2 } ( N - 1 )$ となります．
$K = 2$ でも同様にして，$N - 2$ 個を同色の隣接を許さずに塗ってからブロック $2$ 個を膨らませると考えることができます．$K \geq 2$ では，同じブロックを複数回膨らませて同色が $3$ 個以上する場合も作ることができるという点に注意します．膨らませるブロックを一気に選ぶのではなく，一回ずつどのブロックを膨らませるかを選ぶとすると*3見通しがよくなるかもしれませんが，膨らませる部分は $N - 2$ 箇所から $2$ 回，重複を許して選ぶので，重複組合せ $\left( \dbinom{ N - 2 } 2 \right)$ となっています．同色隣接禁止版の部分の式と合わせれば，$M ( M - 1 )^{ N - 3 }\left( \dbinom { N - 2 } 2 \right)$ となります．
パターンが見えてきました．同色を隣接させる箇所の個数を $i$ とすれば，その場合の塗り方の数は
\begin{align*}
M ( M - 1 )^{ N - 1 - i } \left( \binom{ N - i } i \right) &= M ( M - 1 )^{ N - 1 - i } \binom{ ( N - i ) + i - 1 } i \\
&= M ( M - 1 )^{ N - 1 - i } \binom{ N - 1 } i
\end{align*}
となります．あとはあり得る全ての $i$ について和をとってあげて，
$$
\sum_{ i = 0 }^K M ( M - 1 )^{ N - 1 - i } \binom{ N - 1 } i
$$
を求めれば答えになります．
途中で剰余類環の上での演算や，冪乗・二項係数を計算していますが，これは蟻本 [1] や他のサイト等に詳しいのでそちらに任せます*4．
コード
constexpr int MOD = 998244353; // a^x を mod で求める // 反復二乗法 // O( log x ) long long mod_pow( long long a, long long x, long long mod ); // p が素数のとき、p を法とする剰余体での逆元を求める // Fermat の小定理を利用 // a^{ p - 1 } \equiv 1 ( mod p ) // a^{ p - 2 } \equiv a^{-1} ( mod p ) int mod_inverse( long long a, long long p ); // 素数を法とする剰余体での n! を求める class modFact; // modfact( max_n, mod ) // ()( n ) // ()( n, e ) // nCr class modComb; // modComb( n, mod ) // ()( n, r ) modComb nCr( 200'000, MOD ); int main() { cin.tie( 0 ); ios::sync_with_stdio( false ); cout << setprecision( 12 ) << fixed; IN( int, N, M, K ); LL res = 0; REP( i, K + 1 ) { ( res += LL( M ) * mod_pow( M - 1, N - i - 1, MOD ) % MOD * nCr( N - 1, i ) % MOD ) %= MOD; } cout << res << endl;
参考文献
[1] 秋葉拓哉, 岩田陽一, and 北川宜稔. プログラミングコンテストチャレンジブック: 問題解決のアルゴリズム活用力とコーディングテクニックを鍛える. マイナビ出版, 2012.
*1:重複組合せの公式と例題（玉，整数解の個数） _ 高校数学の美しい物語
*2:$K = 0$ のケースは，上で書いた改変版そのものです
*3:より正確には，重複禁止版のどのブロックに対応する部分を膨らませるかかを選ぶ
*4:問題固有の話からずれるし
