競プロ典型 90 問：069 - Colorful Blocks 2（★3）

問題

atcoder.jp

挑戦結果

• 挑戦日：2021/10/29
• 結果：解けた
• 時間：15分

考えたこと

• 1ブロック目：Kパターン、 2ブロック目：K-1パターン、 3ブロック目以降：K-2パターン なので、これをかければ良い。modを取れば大きな数字ならないので、速度面も問題ないはず。
• 実装したらTLEだった。
• ループがN回(N乗)だけど、Nは と大きかったことに気がついた。
• を N回のループの計算するのではなく、 →二乗→ →二乗→ とやって、logNの計算量で求めるやつだなと思ったけど、Pythonにはライブラリ化されていた気がしたので、ググった。
• pow(x,y,z) で、 をzであまりをとった値になることがわかったので、これに変更してAC。

公式解説

https://twitter.com/e869120/status/1405659999179722754

解説を読んだふりかえり

• 繰り返し二乗法
• powを使ったけど、自分でも実装できるようになったほうが良いのかもしれない。

ソース

N, K = [int(x) for x  in input().split()]
MOD = 10**9 + 7

if N>=2:
    ans = K * (K-1) * pow(K-2, N-2, MOD)
    ans %= MOD
elif N==1:
    ans = K 
    
print(ans)