このシリーズではE869120さんによって作成された競プロ典型90問をPythonで解いています。
問題
ポイント
N = 1, 2, 3以上の場合に条件を満たすような塗り方が何種類あるか見てみましょう。
N = 1 の場合
答え：K通り K種類の色を使用して1つのブロックを塗る方法と同義です。
N = 2 の場合
答え：K * (K-1) 通り
1つ目のブロックを塗る方法がK通り、次のブロックは1つ目で使用した色以外を使って塗るので
K-1通りとなります。
N = 3以上の場合
答え：K * (K-1) * (K-2) ** (N-2) 通り
3つ目のブロックは1,2つ目で使用した色以外で塗るのでK-2通り、4つ目も2,3つ目以外の色で塗るので同じくK-2通り、5つ目も…となり、一般化すると上記の式で求められます。
以上より、Nの数で場合分けして塗り方を求めればOKです。
実装したコードはこちらになります。
コード
def main():
import sys
sys.setrecursionlimit(10 ** 9)
input = sys.stdin.readline
N, K = map(int, input().split(" "))
mod = 10 ** 9 + 7
if N == 1:
print(K % mod)
else: # N >= 2の場合は下記式で求められる
print(K * (K - 1) * pow(K - 2, N - 2,mod) % mod) # K-2のべき乗は(K-2) ** (N-2) でも取れるが、TLEとなる。nが大きい時、余りを取る時はpowを使った方が早い
if __name__ == '__main__':
main()
