Colorful Blocks [AtCoder Beginner Contest 167 E]

https://atcoder.jp/contests/abc167/tasks/abc167_e

解説

https://atcoder.jp/contests/abc167/submissions/13094693

まず、素数mod上での四則演算とコンビネーションが計算できるようにしよう。
二項係数 mod 素数を高速に計算する方法 [累積和, フェルマーの小定理, 繰り返し二乗法, コンビネーション, 10⁹+7] - はまやんはまやんはまやん
このライブラリがあることがACの必要条件。

どこから数え上げようかと思うところであるが、何かを固定することで計算しやすくなる部分が無いか探す。
問題を見ると、固定する部分の候補として「隣り合うブロックで同色である部分の個数」がある。
この個数をk個としておく。
すると、k=0,1,...,Kの場合の組合せの総和を取れば答えになる。

少しシンプルになった。
隣り合うブロックの組であって、同じ色で塗られている組がk個であるとする。
ここからちょっと飛躍が必要なのだが、隣り合うブロック部分はN-1個あるが、
そのN-1個のうち、どこにk個等しくなる部分があるかを考える。
これはC(N-1, k)通りある。
そして、実は、その全ての組について色の塗り方は等しくなり、M×(M-1)^N-1-k通りとなる。

分かりやすさのために、N-1個のうち最初のk個が同色になっていると考えよう。
すると、1番目はM色の選択ができ、そこから次のk個は色が同じである必要があるため、1通り、
それ以降は、前のブロックとは違う色にする必要があるので、M-1通りとなる。
よって、(1番目)×(同じにする)×(異なるようにする)で、M×1^K×(M-1)^N-1-kであり、M×(M-1)^N-1-kとなる。
これが、最初のk個が同色以外のパターンでも考えてみると、前のブロックと同じにする回数と
前のブロックと異なるようにする回数は同じになるため、結局M×(M-1)^N-1-kとなる。

よって、各kについて、C(N-1, k)×M×(M-1)^N-1-kとなり、これの総和を取ると答えになる。
自分の解法では、なんとなくN-1,M-1があるので、N=1とM=1の時は念のため場合分けしている。
もしかしたら、いらないかも。

int N, M, K;
Comb<mint, 1010101> com;
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> M >> K;

    if (N == 1) {
        cout << M << endl;
        return;
    }

    if (M == 1) {
        if (N - 1 == K) cout << M << endl;
        else cout << 0 << endl;
        return;
    }

    mint ans = 0;
    rep(k, 0, K + 1) ans += com.aCb(N - 1, k) * M * (mint(M - 1) ^ (N - 1 - k));
    cout << ans << endl;
}