abc 124 d 'handstand' (python)

直観的に一回の動作で、１と１との間にある０を全て１にする、
次の動作では、前１にしたところより１つ右にある０の塊を全て１にする…
という動作をｋ回繰り返したものの中から連続して１が並ぶ最大値が作れるとわかる
Sのなかに０が連続していくつあるか、１が連続していくつあるかのリストを作り、
2つのリストを足すことで一回の動作で連続する１の個数が作れる
あとはこのリスト a を a[ i ] から a[ i+k-1 ] まで足せばいいのだが、毎回 sum をとっていると TLE するので、はじめにリストを累積和 （a[ i ] を a[ 0 ] ~ a[ i-1 ] の総和）にしておいて、ｋ離れた a[ i ] を引いて求めた
累積和のリストが欲しいときは numpy をつかって np.cumsum(ab,out=ab) とするのが便利
最初が 0 で始まると決めてリストを作った（最初が１から始まる時はその前に０が'０個'並んでいるとしてまず０のリストに'０'を追加した）

def main():
  import sys
  input = sys.stdin.readline
  import numpy as np
  n,k = map(int,input().split())
  s = input()
  chk_0 = []
  chk_1 = []
  c0 = 0
  c1 = 0
  if s[0]=='0':
    c0 += 1
  else:
    c1 += 1
    chk_0.append(0)

  for i in range(n-1):
    if s[i]=='0' and s[i+1]=='0':
      c0 += 1
    if s[i]=='0' and s[i+1]=='1':
      chk_0.append(c0)
      c0 = 0
      c1 = 1
    if s[i]=='1' and s[i+1]=='0':
      chk_1.append(c1)
      c0 = 1
      c1 = 0
    if s[i]=='1' and s[i+1]=='1':
      c1 += 1

  if c0 != 0:
    chk_0.append(c0)
    chk_1.append(0)
  elif c1 != 0:
    chk_1.append(c1)
    chk_0.append(0)

  if len(chk_0) <= k:
    print(n)
    exit()

  ab = np.array([x+y for x,y in zip(chk_0,chk_1)])
  np.cumsum(ab,out=ab)

  ans = ab[k-1]  # 最初はSの左端からｋ個の０の塊を１にしたときの値にしている
  for i in range(len(ab)-k):
    chk = ab[k+i]-ab[i]
    chk += chk_1[i]  # 両端が１になったほうが大きいので、chk の左端に１つ左隣りの１の塊を足す
    ans = max(ans,chk)
  print(ans)
if __name__ == '__main__':
    main()