Atcoder典型90問 069★3[繰り返し２乗法]

069 - Colorful Blocks 2（★3）

■要約

N=1以外の時は，総数がK*(K-1)*(K-2)^(N-2)になることはすぐに分かったが，累乗の計算量はO(N)なのでどうしようかと考えていた．調べたら，繰り返し2乗法というものがあり，計算量をO(logN)に減らせる．

説明は他の方々がとてもわかり易いものを記載しているので割愛する．

pythonだとわざわざ実装しなくても

pow(x,y,n)でx^yをnで割ったあまりをO(logN)で計算してくれるらしい．鬼便利

N, K = map(int, input().split())

mod = 10**9 + 7

if N == 1:
 print(K)
else:
 ans = K*(K-1)*pow(K-2,N-2,mod)
 print(ans%mod)