[AtCoder][典型69] 繰り返し二乗法を使った効率的な累乗の計算

概要

こちらの問題を使用して解説
典型問題69

問題の解き方として、
最初のブロックの選び方はK通り
次のブロックの選び方はK-1通り
それ以降のブロックの選び方はK-2通りなので次の式で解答が出せる。

K * (K-1) * (K-2)^(N-2)

ここで問題になるのは(K-2)^(N-2)の計算量

繰り返し二乗法とは

3^50を考える。
これは3^25 * 3^25に分解することが可能。
同じように3^25は3 * 3^12 * 3^12に分解が可能。
これを繰り返すことで効率良く累乗の計算ができる。

実装

const int MOD = 1e9 + 7;

int pow(long long a, long long n) {
    if (n == 0) return 1;
    if (n == 1) return a % MOD;

    long long ret = pow(a, n / 2) % MOD;
    (ret *= ret) %= MOD;
    if (n % 2 == 1) {
        (ret *= a) %= MOD;
    }
    return ret;
}

int main() {
    long long N, K; cin >> N >> K;
    long long ans = K % MOD;
    if (N >= 2) (ans *= K - 1) %= MOD;
    if (N >= 3) (ans *= pow(K-2, N-2)) %= MOD;
    cout << ans << endl;
}