AtCoder Beginner Contest 124 D 解説動画（尺取法、累積和）をPythonでやってみる

# 概要

2度目のAtCoder Beginner Contestで、今回もD問題が解けず悔しい思いをしました。
今回は解説PDFで実装が省略されていたこともあり、C++での解説動画をPythonでやってみることを通じ、アルゴリズムを理解しようということでやってみました。。
尚、解説動画は生配信で色々とアクシデントがあり面白いので、この記事は見なくても、解説動画を見て、いいねボタンを押しましょう。

[AtCoder Beginner Contest 124 D - Handstand](https://atcoder.jp/contests/abc124/tasks/abc124_d)
[AtCoder Beginner Contest 124 解説配信](https://youtu.be/FRzpDCx17vw)
[AtCoder Beginner Contest 124 解説PDF](https://img.atcoder.jp/abc124/editorial.pdf)

# TLE (Time Limit Exceeded)になるやり方

直大さん(@chokudai)くらいのスーパープログラマーがキレイにコードを書くと通ってしまうことがありますが、普通は通らない実装方法です。

``` python
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)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271789/94195b1f-41c9-426d-326b-16a8b1d37ce9.png)

for文が2重になっていることから計算回数が$O(N^2)$になるため、PythonではしっかりTLEになりました。
個人的には now = 1 とすることで、0から始まる時にNums[0]に0を入れられるところ、now = 1 - now で0と1を交互にできるところ、偶数番目だけを見るなどで、なるほど！となってしまいました。
だいたい末端処理の実装でやらかして、デバックで疲弊してしまうので、こういうテクニックを自然に使えるようになりたいです。

# 尺取法

``` python
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)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271789/77f11527-f399-73b2-b00d-5fe970408a81.png)

尺取法というのは実は解説動画で初めて聞いたのですが、右側と左側の動きが尺取虫のように動くということでしょうね。
計算回数が$O(N)$になるので、Pythonでも余裕をもって通ります。C++では5msということで、差は大きいですね
理論は分かっても実装となると苦労しそうなのですが、left, right, Nextleft, Nextrightという使い方が分かりやすく、メンテナンス性も高いコードで、こういうコードを書きたいですね。

# 累積和

``` python
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)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271789/16ea1b77-42aa-3387-435b-b8e509299d45.png)

累積和という言葉は聞いたことはありましたが、実装はしたことが無いと思います。
計算回数が$O(N)$になり尺取法と同程度ということですが、実行時間が少しだけ尺取法より早くなりました。c++では同じ5msでしたね。
実装は非常に分かりやすいコードで、直大さんも累積和を推奨しますと言っていたので、次に累積和を使う問題が出たらスムーズな実装を目指したいと思います。