ABC167 E - Colorful Blocks - あっとのTECH LOG

問題原文

atcoder.jp

解法

ポイントはまとめて全パターン考えないようにすることだと思う。

まず、「隣り合う色が同じようなものがない」パターンを考えてみる。
これは、例えば一番左は $M$ 色使えて、次からは直前に使った色でない $M- 1$ 色を使えるので、 $M × (M - 1)^{(N - 1)}$ 通りとなる。

この色がバラバラな状態からイメージを膨らませて、「隣り合う色が同じような場所を作る」とは、「２〜 $N$ 番目のブロックのうち、左隣と色が違うものを左隣の色にする」 ことだと考える。

次に、隣り合うブロックの組で色が異なるものが、 $K$ 以下でなく、 $k$ 個として固定して考える。
とすると、通り数は、

(最初の１つで $M$ 通り) × ( $N - 1$ 個の色を揃えられるポイントのうちどの $k$ 個を使うか) × (それぞれのブロックを何色で塗るか)

で求められることになる。

「 $N - 1$ 個の色を揃えられるポイントのうちどの $k$ 個を使うか」については、 ${}_{N - 1} C_k$ で求められる。
「それぞれのブロックを何色で塗るか」については、 $k$ 個分自由に選べる場所が減っているので、 $(M - 1)^{(N - 1 - k)}$ になる。

上記の計算を全ての $k$ に対してやればおしまい！

実装

逆元コンビネーションを使う。

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

ans = 0
for k in range(K + 1):
    ans += M * nCr(N - 1, k) * pow(M - 1, N - 1 - k, MOD) % MOD
    ans %= MOD

print(ans % MOD)

感想

$k$ の結果を $k - 1$ から計算しようとしてダメだった（する必要なかった）
あと 101 みたいなのがあった時に、真ん中を 1 にしたら同じ色の箇所が２箇所増えるじゃんと思っていた。
これは、先にどの部分を同じ色にするかを決めておいて、後から適切な色を割り当てていく とイメージすれば回避できたはず。

上式はぱっと浮かんだのに、余計なことを考えてしまってハマってしまった。。。数え上げ苦手ぇ