概要
2度目のAtCoder Beginner Contestで、今回もD問題が解けず悔しい思いをしました。
今回は解説PDFで実装が省略されていたこともあり、C++での解説動画をPythonでやってみることを通じ、アルゴリズムを理解しようということでやってみました。。
尚、解説動画は生配信で色々とアクシデントがあり面白いので、この記事は見なくても、解説動画を見て、いいねボタンを押しましょう。
AtCoder Beginner Contest 124 D - Handstand
AtCoder Beginner Contest 124 解説配信
AtCoder Beginner Contest 124 解説PDF
TLE (Time Limit Exceeded)になるやり方
直大さん(@chokudai)くらいのスーパープログラマーがキレイにコードを書くと通ってしまうことがありますが、普通は通らない実装方法です。
N, K = map(int, input().split())
S = input()
Nums = []
now = 1 # 今見ている数
cnt = 0 # nowがいくつ並んでいるか
for i in range(N):
if S[i] == str(now):
cnt += 1
else:
Nums.append(cnt)
now = 1 - now # 0と1を切り替えるときの計算 now ^= 1
cnt = 1
if cnt != 0:
Nums.append(cnt)
# 1-0-1-0-1-0-1 って感じの配列が欲しい
# 1-0-1-0-1-0 みたいに0で終わっていたら適当に1つ足す
if len(Nums) % 2 == 0:
Nums.append(0)
Add = 2 * K + 1
ans = 0
# 1-0-1... の1から始めるので、偶数番目だけ見る
for i in range(0, len(Nums), 2):
tmp = 0
left = i
right = min(i + Add, len(Nums))
for j in range(left, right):
tmp += Nums[j]
ans = max(tmp, ans)
print(ans)
for文が2重になっていることから計算回数が$O(N^2)$になるため、PythonではしっかりTLEになりました。
個人的には now = 1 とすることで、0から始まる時にNums[0]に0を入れられるところ、now = 1 - now で0と1を交互にできるところ、偶数番目だけを見るなどで、なるほど！となってしまいました。
だいたい末端処理の実装でやらかして、デバックで疲弊してしまうので、こういうテクニックを自然に使えるようになりたいです。
尺取法
N, K = map(int, input().split())
S = input()
Nums = []
now = 1 # 今見ている数
cnt = 0 # nowがいくつ並んでいるか
for i in range(N):
if S[i] == str(now):
cnt += 1
else:
Nums.append(cnt)
now = 1 - now # 0と1を切り替えるときの計算 now ^= 1
cnt = 1
if cnt != 0:
Nums.append(cnt)
# 1-0-1-0-1-0-1 って感じの配列が欲しい
# 1-0-1-0-1-0 みたいに0で終わっていたら適当に1つ足す
if len(Nums) % 2 == 0:
Nums.append(0)
Add = 2 * K + 1
ans = 0
# 尺取り法 forループの外側にleft, rightを持つ
left = 0
right = 0
tmp = 0 # [left, right)のsum
# 1-0-1... の1から始めるので、偶数番目だけ見る
for i in range(0, len(Nums), 2):
# 次のleft, rightを計算する
Nextleft = i
Nextright = min(i + Add, len(Nums))
# 左端を移動する
while Nextleft > left:
tmp -= Nums[left]
left += 1
# 右端を移動する
while Nextright > right:
tmp += Nums[right]
right += 1
ans = max(tmp, ans)
print(ans)
尺取法というのは実は解説動画で初めて聞いたのですが、右側と左側の動きが尺取虫のように動くということでしょうね。
計算回数が$O(N)$になるので、Pythonでも余裕をもって通ります。C++では5msということで、差は大きいですね
理論は分かっても実装となると苦労しそうなのですが、left, right, Nextleft, Nextrightという使い方が分かりやすく、メンテナンス性も高いコードで、こういうコードを書きたいですね。
累積和
N, K = map(int, input().split())
S = input()
Nums = []
now = 1 # 今見ている数
cnt = 0 # nowがいくつ並んでいるか
for i in range(N):
if S[i] == str(now):
cnt += 1
else:
Nums.append(cnt)
now = 1 - now # 0と1を切り替えるときの計算 now ^= 1
cnt = 1
if cnt != 0:
Nums.append(cnt)
# 1-0-1-0-1-0-1 って感じの配列が欲しい
# 1-0-1-0-1-0 みたいに0で終わっていたら適当に1つ足す
if len(Nums) % 2 == 0:
Nums.append(0)
Add = 2 * K + 1
# 累積和を作る
# 0 1 2 3 4 5 6
# 0 1 2 3 4 5
sum = [0] * (len(Nums) + 1)
for i in range(len(Nums)):
sum[i + 1] = sum[i] + Nums[i]
ans = 0
# 1-0-1... の1から始めるので、偶数番目だけ見る
for i in range(0, len(Nums), 2):
# 次のleft, rightを計算する [left, right)
left = i
right = min(i + Add, len(Nums))
tmp = sum[right] - sum[left]
ans = max(tmp, ans)
print(ans)
累積和という言葉は聞いたことはありましたが、実装はしたことが無いと思います。
計算回数が$O(N)$になり尺取法と同程度ということですが、実行時間が少しだけ尺取法より早くなりました。c++では同じ5msでしたね。
実装は非常に分かりやすいコードで、直大さんも累積和を推奨しますと言っていたので、次に累積和を使う問題が出たらスムーズな実装を目指したいと思います。