AtCoder の問題を解くために必要な実力を付けるために作られた「典型問題」を解いていく企画として、「典型90」があります。
今回は、その内、69の問目である「069 – Colorful Blocks 2（★3）」を解説していきたいと思います。
なにか、間違っているところなどありましたら、@EitaSugiまで、ご指導ご鞭撻の程、よろしくお願い致します！
問題名 【069 – Colorful Blocks 2（★3）】
問題概要
N 個のブロックが横一列に並んでおり、左から順に 1 から N までの番号が付けられています。これらのブロックそれぞれを、K 種類の色のうちいずれか一色で塗ることを考えます。ただし、次の条件を満たすように塗らなければなりません。
- 1≤∣i−j∣≤2 ならば、ブロック i とブロック j に塗られている色は異なる。
このようなブロックの塗り方が何通りあるかを、 \(10^9+7\) で割った余りを出力する問題です。
考察
NやKの制約が大きいため、全探索は不可能です。しかし、問題の性質をよく観察すると、ある規則性が見えてきます。
- N=1のとき、K通りの塗り方があります。
- N=2のとき、最初のブロックにK通り、次のブロックにK-1通りの塗り方があります。
- N>=3のとき、最初の3つのブロックの塗り方は \(K×(K−1)×(K−2)\) 通りです。そして、4番目以降のブロックについては、前のブロックの色に依存して、\(K-2\)通りの塗り方が増えていきます。
この規則性を利用して、繰り返し二乗法を使って高速に計算を行うことができます。
典型解法：繰り返し二乗法
コード
N,K = map(int, input().split())
mod = 10**9 + 7
if N == 1:
print(K)
exit()
if K == 1:
print(0)
exit()
if N == 2:
print(K*(K-1))
exit()
if K == 2:
print(0)
exit()
# 繰り返し二乗法で (K - 2) ** (N - 3) を mod で割った余りを計算
def mod_pow(x, n, mod):
result = 1
while n > 0:
if n % 2 == 1:
result = result * x % mod
x = x * x % mod
n //= 2
return result
# 本体の計算
ans = mod_pow(K - 2, N - 3, mod)
for i in range(3):
ans = ans * (K - i) % mod
print(ans)まとめ
この問題では、ブロックの塗り方の総数を求めるための方法を考察し、その考察を元に実装したコードを紹介しました。繰り返し二乗法を利用することで、大きな制約の中でも高速に計算を行うことができました。
コメント