AtCoder Beginner Contest 167 E Colorful Blocks

$N$ 個のブロックがあり、これを $M$ 色のペンキで塗り分ける。隣り合う同色のブロックが $K$ 組以内の場合の塗り分け方の総数を $mod\ 998244353$ で求めよ。

左端を除くブロックのうち隣と同じ色で塗られているブロックの数が丁度 $k$ 個の場合を考える。

この場合の色の塗り方は以下を考慮して計算できる。

1) 左隣と同じ色を持つブロックを $k$ 個選ぶ => $_{N-1}C_{k}$

2) 1)で選択したブロックを除く $N-1-k$ のブロックを左隣と違う色で塗る => $(M-1) ^{ N-1-k }$

3) 左端のブロックは $M$ 色から1つ選択して塗る => $M$

よってこの場合の色の塗り方は $_{N-1}C_{k} * (M-1) ^ {N-1-k} * M$

後は $k = 0$ から $k = K$ までの塗り方の和を上の式に当てはめて計算すれば良い。

この提出では以下の記事で紹介されたライブラリを参考にしました。

seriruの技術屋ブログ