Harryxun/testdata huggingface dataset:
Buggy python submission (verdict: Runtime Error):
def COMinit():
    fac[0] = fac[1] = 1
    finv[0] = finv[1] = 1
    inv[1] = 1
    for i in range(2,max):
        fac[i] = fac[i-1]*i%mod
        inv[i] = mod - inv[mod%i]*(mod//i)%mod
        finv[i] = finv[i-1]*inv[i]%mod

def COM(n,k):
    if n < k:
        return 0
    if n < 0 or k < 0:
        return 0
    return fac[n] * (finv[k]*finv[n-k]%mod)%mod

mod = 998244353
N,M,K = map(int,input().split())

max = N 
fac = [0] * max
finv = [0] * max
inv = [0] * max
COMinit()
ans = 0
for i in range(K+1):
    c = COM(N-1,i)
    m = pow(M-1,N-i-1,mod)
    ans = (ans + ((c*m%mod)*M%mod))%mod
print(ans)


Fixed submission:
def COMinit():
    fac[0] = fac[1] = 1
    finv[0] = finv[1] = 1
    inv[1] = 1
    for i in range(2,max):
        fac[i] = fac[i-1]*i%mod
        inv[i] = mod - inv[mod%i]*(mod//i)%mod
        finv[i] = finv[i-1]*inv[i]%mod

def COM(n,k):
    if n < k:
        return 0
    if n < 0 or k < 0:
        return 0
    return fac[n] * (finv[k]*finv[n-k]%mod)%mod

mod = 998244353
N,M,K = map(int,input().split())

max = N + M + K
fac = [0] * max
finv = [0] * max
inv = [0] * max
COMinit()
ans = 0
for i in range(K+1):
    c = COM(N-1,i)
    m = pow(M-1,N-i-1,mod)
    ans = (ans + ((c*m%mod)*M%mod))%mod
print(ans)

Buggy python submission (verdict: Runtime Error):
def COMinit():
    fac[0] = fac[1] = 1
    finv[0] = finv[1] = 1
    inv[1] = 1
    for i in range(2,max):
        fac[i] = fac[i-1]*i%mod
        inv[i] = mod - inv[mod%i]*(mod//i)%mod
        finv[i] = finv[i-1]*inv[i]%mod

def COM(n,k):
    if n < k:
        return 0
    if n < 0 or k < 0:
        return 0
    return fac[n] * (finv[k]*finv[n-k]%mod)%mod

mod = 998244353
N,M,K = map(int,input().split())

max = N
fac = [0] * max
finv = [0] * max
inv = [0] * max
COMinit()
ans = 0
for i in range(K+1):
    c = COM(N-1,i)
    m = pow(M-1,N-i-1,mod)
    ans = (ans + ((c*m%mod)*M%mod))%mod
print(ans)


Fixed submission:
def COMinit():
    fac[0] = fac[1] = 1
    finv[0] = finv[1] = 1
    inv[1] = 1
    for i in range(2,max):
        fac[i] = fac[i-1]*i%mod
        inv[i] = mod - inv[mod%i]*(mod//i)%mod
        finv[i] = finv[i-1]*inv[i]%mod

def COM(n,k):
    if n < k:
        return 0
    if n < 0 or k < 0:
        return 0
    return fac[n] * (finv[k]*finv[n-k]%mod)%mod

mod = 998244353
N,M,K = map(int,input().split())

max = N + M + K
fac = [0] * max
finv = [0] * max
inv = [0] * max
COMinit()
ans = 0
for i in range(K+1):
    c = COM(N-1,i)
    m = pow(M-1,N-i-1,mod)
    ans = (ans + ((c*m%mod)*M%mod))%mod
print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
S = input()
ans = 0
for i in range(len(S)):
    for j in range(i+4,len(S)):
        if int(S[i:j+1]) % 2019 == 0:

            ans += 1
print(ans)
#3*673


Fixed submission:
S = input()
mod = 2019

array = []
for i in range(len(S)):
    x = (int(S[len(S)-1-i])*pow(10,i,mod))%mod
    array.append(x)
array2 = [0]
y = 0
for i in range(len(S)):
    y = (y+array[i])%mod
    array2.append(y)
array3 = [0] * 2019
ans = 0
for i in range(len(array2)):
    z = array2[i]
    ans += array3[z]
    array3[z] += 1
print(ans)
#3*673
Buggy python submission (verdict: Runtime Error):
from collections import deque

def bfs():
    # たどり着くまでにかかった回数をいれる多重リスト。
    d = [[float("inf")] * w for i in range(h)]
    # 1ノード降りるときにforで回せるようリスト化しておく
    dx = [1, 0, -1, 0]
    dy = [0, 1, 0, -1]
    # 接したポイントを取り込む用のqueを定義
    que = deque([])
    # スタート地点をqueに加える。探索の際はqueから取り出すループを回すので最初だけ先に入れておく
    que.append(sx,sy)
    # スタートからスタートにかかる距離は0なのでinfを0にかえる。
    d[sx][sy] = 0
    # queがあるかぎり続ける
    while que:
        p = que.popleft()
        # 取り出した点がゴールなら終了
        if p[0] == gy and p[1] == gx:
            break
        #　それ以外の時は，接している点を4方向順に取得。一つずつ処理する。
        for i in range(4):
            nx = p[0] + dx[i]
            ny = p[1] + dy[i]
            # 新しく取得した座標に対しての処理
            # 迷路内に存在する座標であり，壁ではなく，まだ通ったことがなくdがinfの点なら処理する = 新しく取得した道候補の1つ
            if 0 <= nx < r and 0 <= ny < c and maze[nx][ny] != "#" and d[nx][ny] == float("inf"):
                # 道候補をqueに加える。
                que.append((nx, ny))
                # その道候補への経路までにかかった経路をdに加える。
                d[nx][ny] = d[p[0]][p[1]] + 1
    return d[gy][gx]

h, w = map(int, input().split())
maze = [list(input()) for i in range(h)]
sx, sy = 0, 0
gx, gy = h - 1, w - 1

white = 0
for i in range(h):
    for j in range(w):
        if maze[i][j] == ".":
            white += 1

res = bfs()
if 0 < res < float("inf"):
    # 白いマスの数から最短経路でかかるコスト分を引く
    # ゴールを黒いマスに変えることはできないためその分も引く
    print(white - res - 1)
else:
    print(-1)


Fixed submission:
from collections import deque

def bfs():
    # たどり着くまでにかかった回数をいれる多重リスト。
    d = [[float("inf")] * w for i in range(h)]
    # 1ノード降りるときにforで回せるようリスト化しておく
    dx = [1, 0, -1, 0]
    dy = [0, 1, 0, -1]
    # 接したポイントを取り込む用のqueを定義
    que = deque([])
    # スタート地点をqueに加える。探索の際はqueから取り出すループを回すので最初だけ先に入れておく
    que.append((sx,sy))
    # スタートからスタートにかかる距離は0なのでinfを0にかえる。
    d[sx][sy] = 0
    # queがあるかぎり続ける
    while que:
        p = que.popleft()
        # 取り出した点がゴールなら終了
        if p[0] == gy and p[1] == gx:
            break
        #　それ以外の時は，接している点を4方向順に取得。一つずつ処理する。
        for i in range(4):
            nx = p[0] + dx[i]
            ny = p[1] + dy[i]
            # 新しく取得した座標に対しての処理
            # 迷路内に存在する座標であり，壁ではなく，まだ通ったことがなくdがinfの点なら処理する = 新しく取得した道候補の1つ
            if 0 <= nx < h and 0 <= ny < w and maze[nx][ny] != "#" and d[nx][ny] == float("inf"):
                # 道候補をqueに加える。
                que.append((nx, ny))
                # その道候補への経路までにかかった経路をdに加える。
                d[nx][ny] = d[p[0]][p[1]] + 1
    return d[gx][gy]

h, w = map(int, input().split())
maze = [list(input()) for i in range(h)]
sx, sy = 0, 0
gx, gy = h - 1, w - 1

white = 0
for i in range(h):
    for j in range(w):
        if maze[i][j] == ".":
            white += 1

res = bfs()
if 0 < res < float("inf"):
    # 白いマスの数から最短経路でかかるコスト分を引く
    # ゴールを黒いマスに変えることはできないためその分も引く
    print(white - res - 1)
else:
    print(-1)
Buggy python submission (verdict: Runtime Error):
try:
	n = int(input())
	tmp = 360
	while True:
    	if tmp % n == 0:
        	print(tmp//n)
        	exit(0)
    	else:
        	tmp+= 360

Fixed submission:
X = int(input())
tmp = 360
while True:
    if tmp % X == 0:
        print(tmp//X)
        exit()
    else:
        tmp+= 360
Buggy python submission (verdict: Time Limit Exceeded):
#!/usr/bin/env python3

from itertools import combinations

N = int(input())
A = map(int, input().split())

total = sum([x * y for x, y in list(combinations(A, 2))])
result = total % (10**9 + 7)
print(result)


Fixed submission:
#!/usr/bin/env python3


def solve(N, A):
    MOD = 10**9 + 7
    total = sum(A)
    ans = 0

    for i in range(N - 1):
        total -= A[i]
        ans += A[i] * total
        ans %= MOD

    return ans


def main():
    N = int(input())
    A = list(map(int, input().split()))

    ans = solve(N, A)

    print(ans)


if __name__ == "__main__":
    main()
Buggy python submission (verdict: Runtime Error):
#!/usr/bin/env python3

X, Y = map(int, input().split())

# a + b = X
# 2a + 4b = Y
# --
# b = (Y-2X) / 2

b = (Y - 2 * X) / 2.0
a = X - b

if a.is_integer() and b.is_integer() and a >= 0 and b >= 0:
    print("Yes")
else:
    print("No")


Fixed submission:
#!/usr/bin/env python3

a = list(map(int, input().split()))
print(a.index(0) + 1)
Buggy python submission (verdict: Time Limit Exceeded):
#!/usr/bin/env python3

N = int(input())
A = list(map(int, input().split()))
A.sort()

count = 0

for i in range(N):
    for j in range(N):
        if i == j:
            continue

        if A[i] % A[j] == 0:
            break

    else:
        count += 1

print(count)


Fixed submission:
#!/usr/bin/env python3

N = int(input())
A = list(map(int, input().split()))

A.sort()
A_max = A[-1]
table = [0] * (A_max + 1)

for x in A:
    table[x] += 1
    for y in range(x * 2, A_max + 1, x):
        table[y] += 1

result = sum([1 for x in A if table[x] == 1])
print(result)

Buggy python submission (verdict: Time Limit Exceeded):
#!/usr/bin/env python3

N = int(input())
A = list(map(int, input().split()))

count = 0

for i in range(N):
    for j in range(N):
        if i == j:
            continue

        if A[i] % A[j] == 0:
            break

    else:
        count += 1

print(count)


Fixed submission:
#!/usr/bin/env python3

N = int(input())
A = list(map(int, input().split()))

A.sort()
A_max = A[-1]
table = [0] * (A_max + 1)

for x in A:
    table[x] += 1
    for y in range(x * 2, A_max + 1, x):
        table[y] += 1

result = sum([1 for x in A if table[x] == 1])
print(result)
Buggy python submission (verdict: Runtime Error):
#!/usr/bin/env python3

N = int(input())
l = list(range(2, N + 1)) + [1]
print(sum([i % x for i, x in enumerate(l, start=1)]))


Fixed submission:
#!/usr/bin/env python3

N = int(input())
print(N * (N - 1) // 2)
Buggy python submission (verdict: Time Limit Exceeded):
#!/usr/bin/env python3
import sys
read = sys.stdin.buffer.read

K, S = map(int, read().split())
count = 0

for x in range(K + 1):
    for y in range(K + 1):
        for z in range(K + 1):
            if x + y + z == S:
                count += 1

print(count)


Fixed submission:
#!/usr/bin/env python3
import sys
read = sys.stdin.buffer.read

K, S = map(int, read().split())
count = 0

for x in range(K + 1):
    for y in range(K + 1):
        z = S - x - y
        if 0 <= z <= K:
            count += 1

print(count)
Buggy python submission (verdict: Runtime Error):
# C - Forbidden List

x, n = map(int,input().split())
P = list(map(int,input().split()))

t = -100
a = [0]*100
for p in P:
    a[p-1] = 1

d = 0
while True:
    if a[x-1-d] == 0:
        t = x-d
        break
    elif a[x-1+d] == 0:
        t = x + d
        break
    else:
        d += 1

print(t)


Fixed submission:
# C - Forbidden List

x, n = map(int,input().split())
P = list(map(int,input().split()))

t = -100
a = [0]*300
for p in P:
    a[100+p-1] = 1

d = 0
while True:
    if a[100+x-1-d] == 0:
        t = x-d
        break
    elif a[100+x-1+d] == 0:
        t = x + d
        break
    else:
        d += 1

print(t)
Buggy python submission (verdict: Time Limit Exceeded):
# C - gacha

n = int(input())
s = []

for i in range(n):
    gacha = input()
    if gacha not in s:
        s.append(gacha)

print(len(s))


Fixed submission:
# C - gacha

n = int(input())
s = []
c = 1

for i in range(n):
    s.append(input())

s.sort()

for j in range(1,n):
    if s[j-1] != s[j]:
        c += 1


print(c)

Buggy python submission (verdict: Time Limit Exceeded):
# C - gacha

n = int(input())
s = []

for i in range(n):
    gacha = input()
    if gacha not in s:
        s.append(gacha)

print(len(s))


Fixed submission:
# C - gacha

n = int(input())
s = []
c = 1

for i in range(n):
    s.append(input())

s.sort()

for j in range(1,n):
    if s[j-1] != s[j]:
        c += 1


print(c)
Buggy python submission (verdict: Time Limit Exceeded):
# C - Sum of gcd of Tuples (Easy)
import math

k = int(input())
s = 0

for a in range(1,k+1):
    for b in range(1,k+1):
        for c in range(1,k+1):
            s += math.gcd(a,math.gcd(b,c))

print(s)


Fixed submission:
# C - Sum of gcd of Tuples (Easy)

import math

k = int(input())
s = 0
for a in range(1,k+1):
    for b in range(1,k+1):
        gcdab = math.gcd(a,b)
        if  gcdab== 1:
            s += k
        else:
            for c in range(1,k+1):
                s += math.gcd(gcdab,c)

print(s)
Buggy python submission (verdict: Time Limit Exceeded):
# C - Fennec vs Monster

n,k = map(int,input().split())
h = list(map(int,input().split()))
h.sort(reverse=True)
if k >= n:
    k = n
for i in range(k):
    del h[0]

print(sum(h))


Fixed submission:
# C - Fennec vs Monster

n,k = map(int,input().split())
h = list(map(int,input().split()))
h.sort(reverse=True)
if k >= n:
    k = n
h2 = h[k:]

print(sum(h2))
Buggy python submission (verdict: Time Limit Exceeded):
n,k=[int(c) for c in input().split()]
MOD = 998244353
segs = set()
for _ in range(k):
    a,b=[int(c) for c in input().split()]
    for i in range(a,b+1):
        segs.add(i)

dp =[0 for i in range(n)]
dp[0] = 1
# print(segs)

for i in range(n-1):
    if dp[i] != 0:
        for j in segs:
            if i + j < n:
                dp[i+j] = dp[i+j] + dp[i]
                dp[i+j]%=MOD

# print(dp)
print(dp[-1])



Fixed submission:
n,k=[int(c) for c in input().split()]
MOD = 998244353
segs = []
for _ in range(k):
    a,b=[int(c) for c in input().split()]
    segs.append((a,b))


dp=[0 for i in range(n+5)]
dp[0] = 1

for i in range(n):
    # dp[i] = dp[i] + dp[i-1]
    if dp[i] < 0:
        dp[i]+=MOD
    for a,b in segs:
        if i+a <=n:
            dp[i+a]+=dp[i]
        if i+b+1 <= n:
            dp[i+b+1]-=dp[i]
    dp[i]%=MOD
    if i !=0:
        dp[i+1] = dp[i+1]+dp[i]

print(dp[n-1]%MOD)
Buggy python submission (verdict: Runtime Error):
def find_par(x,p):
    if x == p[x] :
        return x

    p[x] = find_par(p[x],p)
    return p[x]


def union(x,y,p,size):
    xp = find_par(x,p)
    yp= find_par(y,p)

    if xp != yp:
        p[yp] = xp
        size[xp] += size[yp]
    
    
n,m=[int(c) for c in input().split()]
p=[i for i in range(n+1)]
size =[1 for i in range(n+1)]
for i in range(m):
    a,b =[int(c) for c in input().split()]
    union(a,b,p,size)

roots= set()

for i in range(1,n+1):
    par = find_par(i,p)
    roots.add(par)

ans = float("-inf")
for pp in roots:
    ans = max(ans,size[pp])

print(ans)

Fixed submission:
def find_par(x,p):
    
    try :
        if x == p[x] :
            return x

        p[x] = find_par(p[x],p)
        return p[x]
    except:
        return x


def union(x,y,p,size):
    try:
        xp = find_par(x,p)
        yp= find_par(y,p)

        if xp != yp:
            p[yp] = xp
            size[xp] += size[yp]
    except:
        pass
    
n,m=[int(c) for c in input().split()]
p=[i for i in range(n+1)]
size =[1 for i in range(n+1)]
for i in range(m):
    a,b =[int(c) for c in input().split()]
    union(a,b,p,size)

roots= set()

for i in range(1,n+1):
    par = find_par(i,p)
    roots.add(par)

ans = float("-inf")
for pp in roots:
    try:
        ans = max(ans,size[pp])
    except:
        pass

print(ans)
Buggy python submission (verdict: Runtime Error):
def mpow(a,b,m):
    ans=1
    while b >0 :
        if b&1:
            ans = (ans*a)%m
        a=(a*a)%m
        b=b>>1

    return ans

def calcmod(X,m,N):
    mod=0
    for i in range(N+1):
        if X&(1<<i) > 0:
            mod += mpow(2,i,m)
            mod%=m
    return mod

def popcount(m):
    return bin(m).count("1")


def findsteps(mm):
    cnt=0
    while mm !=0:
        cnt+=1
        mm=mm%popcount(mm)
    return cnt
N=int(input())
X=int(input(),2)

##we need to find X%(m-1) and X%(m+1)
m=popcount(X)
m1=m+1
m2=m-1
firstmod=calcmod(X,m1,N)
if m2 !=0:
    secondmod=calcmod(X,m2,N)
fans=[0 for i in range(N)]
k=0
for i in range(N-1,-1,-1):
    if X & (1<<i) >0:
        ##the bit was set
        ##we need to find X%(m-1) - (2^i)%(m-1)
        if m2 == 0:
            ans=0
        else:
            mm=secondmod - mpow(2,i,m2)
            if mm < 0:
                mm+=m2
            mm=mm%m2
            ans=1+findsteps(mm)
    else:
        mm = firstmod + mpow(2,i,m1)
        mm=mm%m1
        ans=1+findsteps(mm)
    fans[k] = ans
    k+=1
        ##the bit was not set
for i in fans:
    print(i)

Fixed submission:
L,R,d=[int(c) for c in input().split()]
ans=0
for i in range(L,R+1):
    if i%d == 0:
        ans+=1

print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
def main(): 
    n,k = map(int,input().split())
    LR = [list(map(int,input().split())) for _ in range(k)]

    dp = [0]*(n*2 + 100)
    dp[1] = 1
    dp[2] = -1
    for i in range(1,n + 1):
        dp[i] += dp[i - 1]
        for l,r in LR:
            if i + l > n:
                continue
            dp[l + i] += dp[i]
            dp[r + i + 1] -= dp[i]        
    print(dp[n]%998244353) 
    return

if __name__ == "__main__":
    main()

Fixed submission:
def main(): 
    mod = 998244353
    n,k = map(int,input().split())
    LR = [list(map(int,input().split())) for _ in range(k)]

    dp = [0]*(n*2 + 100)
    dp[1] = 1
    dp[2] = -1
    for i in range(1,n + 1):
        dp[i] = (dp[i] + dp[i - 1])%mod
        for l,r in LR:
            if i + l > n:
                continue
            dp[l + i] = (dp[l+i] + dp[i])%mod
            dp[r + i + 1] = dp[r+i+1] - dp[i]  if dp[r+i+1] - dp[i] > 0 else dp[r+i+1] - dp[i] + 998244353
    print(dp[n]%998244353) 
    return

if __name__ == "__main__":
    main()
Buggy python submission (verdict: Runtime Error):
    from math import log10

    class Mod:
        fac = []
        inv = []
        finv = []
        mod = 0
        def __init__(self,maxi:int ,mod:int):
            self.mod = mod
            self.fac = [1]*maxi
            self.inv = [1]*maxi
            self.finv = [1]*maxi
            for i in range(2,maxi):
                self.fac[i] = self.fac[i - 1]*i%mod
                self.inv[i] = mod - self.inv[mod%i]*(mod//i)%mod
                self.finv[i] = self.finv[i-1]*self.inv[i]%mod


        def com(self,n:int, k:int)->int:
            #return nCk
            if n < k:
                return 0
            if n < 0 or k < 0:
                return 0
            return self.fac[n]*(self.finv[k]*self.finv[n-k]%self.mod)%self.mod


    def main():
        s = int(input())
        m =  Mod(s,10**9 + 7)
        ans = 0
        for n in range(1,s // 3 + 1):
            if 3*n == s:
                ans = (ans+1)%(10**9 + 7)
                break
            ans = (ans + m.com(-2*n + s - 1,s - 3*n))%(10**9 + 7)
        print(ans)
        return

    if __name__ == "__main__":
        main()

Fixed submission:
from math import log10

class Mod:
    fac = []
    inv = []
    finv = []
    mod = 0
    def __init__(self,maxi:int ,mod:int):
        self.mod = mod
        self.fac = [1]*maxi
        self.inv = [1]*maxi
        self.finv = [1]*maxi
        for i in range(2,maxi):
            self.fac[i] = self.fac[i - 1]*i%mod
            self.inv[i] = mod - self.inv[mod%i]*(mod//i)%mod
            self.finv[i] = self.finv[i-1]*self.inv[i]%mod


    def com(self,n:int, k:int)->int:
        #return nCk
        if n < k:
            return 0
        if n < 0 or k < 0:
            return 0
        return self.fac[n]*(self.finv[k]*self.finv[n-k]%self.mod)%self.mod


def main():
    s = int(input())
    m =  Mod(s,10**9 + 7)
    ans = 0
    for n in range(1,s // 3 + 1):
        if 3*n == s:
            ans = (ans+1)%(10**9 + 7)
            break
        ans = (ans + m.com(-2*n + s - 1,s - 3*n))%(10**9 + 7)
    print(ans)
    return

if __name__ == "__main__":
    main()
Buggy python submission (verdict: Runtime Error):
from math import ceil

n,k = map(int,input().split())
A = list(map(int,input().split()))

def flag(upper):
    c = 0
    for a in A:
        c += int(a / upper)
        if c > k:
            return False
    return True

def main():
    A.sort(reverse=True)
    low = 0
    high = A[0]
    while low < high:
        mid = (low + high) // 2
        if flag(mid):
            high = mid
        else:
            low = mid + 1
    print(high)
    return

if __name__ == "__main__":
    main()

Fixed submission:
from math import ceil

n,k = map(int,input().split())
A = list(map(int,input().split()))

def flag(upper):
    c = 0
    for a in A:
        c += ceil(a / upper) - 1
        if c > k:
            return False
    return True

def main():
    A.sort(reverse=True)
    low = 1
    high = A[0]
    while low < high:
        mid = (low + high) // 2
        if flag(mid):
            high = mid
        else:
            low = mid + 1
    print(high)
    return

if __name__ == "__main__":
    main()
Buggy python submission (verdict: Runtime Error):
from math import log2

def main():
    N,K = map(int,input().split())
    A = list(map(int,input().split()))
    if K > log2(N):
        return print(*[N]*N,sep = " ")
    for _ in range(K):
        imos = [0]*(N + 1)
        for i in range(N):
            imos[max(0,i  - A[i])] += 1
            imos[min(N,i + 1 + A[i])] -= 1
        for j in range(N):
            imos[j+1] += imos[j]
            flag += 1 if imos[j] == N else 0
        A = imos[:N].copy()
    print(*A,sep=" ")
    return 

if __name__ == "__main__":
    main()

Fixed submission:
from math import log2

def main():
    N,K = map(int,input().split())
    A = list(map(int,input().split()))
    if K > log2(N)+ 50:
        return print(*[N]*N,sep = " ")
    for _ in range(K):
        imos = [0]*(N + 1)
        for i in range(N):
            imos[max(0,i  - A[i])] += 1
            imos[min(N,i + 1 + A[i])] -= 1
        for j in range(N):
            imos[j+1] += imos[j]
        A = imos[:N].copy()
    print(*A,sep=" ")
    return

if __name__ == "__main__":
    main()

Buggy python submission (verdict: Runtime Error):
def main():
    N,K = map(int,input().split())
    A = list(map(int,input().split()))
    flag = 0
    for k in range(K):
        imos = [0]*(N + 1)
        if flag == N:
            print([N]*(N))
            return
        for i in range(N):
            imos[max(0,i  - A[i])] += 1
            imos[min(N + 1,i + 1 + A[i])] -= 1
        for j in range(N):
            imos[j+1] += imos[j]
        A = imos[:N].copy()
    print(*A,sep=" ")
    return 

if __name__ == "__main__":
    main()

Fixed submission:
from math import log2

def main():
    N,K = map(int,input().split())
    A = list(map(int,input().split()))
    if K > log2(N)+ 50:
        return print(*[N]*N,sep = " ")
    for _ in range(K):
        imos = [0]*(N + 1)
        for i in range(N):
            imos[max(0,i  - A[i])] += 1
            imos[min(N,i + 1 + A[i])] -= 1
        for j in range(N):
            imos[j+1] += imos[j]
        A = imos[:N].copy()
    print(*A,sep=" ")
    return

if __name__ == "__main__":
    main()
Buggy python submission (verdict: Time Limit Exceeded):
from math import ceil, floor

# dp[a_i番まで見た。a_iを含む][a_i（を含む）までにいくつ採用しているか][i - 1を使ったかどうか] bool


def main():
    n = int(input())
    A = [0] + list(map(int, input().split()))
    dp = {}
    dp[(0, 0)] = 0
    dp[(1, 1)] = A[1]
    dp[(1, 0)] = 0
    for i in range(2, n + 1):
        for j in range(max(floor(n/2) - ceil((n-i)/2), 0), floor(n/2) + 1):
            # a_iを採用する場合
            if (i, j) in dp:
                dp[(i, j)] = max(dp[(i, j)], A[i] + dp[(i - 2, j - 1)]
                                 ) if (i - 2, j - 1) in dp else dp[(i, j)]
            else:
                dp[(i, j)] = A[i] + dp[(i - 2, j - 1)] if (i -
                                                           2, j - 1) in dp else float("inf")*-1
            # a_iを採用しない場合
            if (i - 1, j) in dp:
                dp[(i, j)] = max(dp[(i - 1, j)], dp[(i, j)])

    print(dp[(n, floor(n/2))])
    return


"""
7
0を取ったところ
#######
#0#0#0#
0#0#0##
0#0##0#
0#0###0
##0#0#0

4
0#0#
#0#0
0##0
"""
if __name__ == "__main__":
    main()


Fixed submission:
from math import ceil, floor

# dp[a_i番まで見た。a_iを含む][a_i（を含む）までにいくつ採用しているか][i - 1を使ったかどうか] bool
memo = {}


def dp(i, j):
    if (i, j) in memo:
        return memo[(i, j)]
    else:
        return float("inf")*-1


def main():
    n = int(input())
    A = [0] + list(map(int, input().split()))
    memo[(0, 0)] = 0
    memo[(1, 1)] = A[1]
    memo[(1, 0)] = 0
    for i in range(2, n + 1):
        for j in range(max(floor(n/2) - ceil((n-i)/2) - 1, 0), floor(i/2) + 5):
            # a_iを採用する場合
            if (i, j) in memo:
                memo[(i, j)] = max(
                    [dp(i, j), A[i] + dp(i - 2, j - 1), A[i] + dp(i - 3, j - 1)])
            else:
                memo[(i, j)] = max(A[i] + dp(i - 2, j - 1),
                                   A[i] + dp(i - 3, j - 1))
            # a_iを採用しない場合
            if (i - 1, j) in memo:
                memo[(i, j)] = max(dp(i - 1, j), dp(i, j))

    print(dp(n, floor(n/2)))
    return
if __name__ == "__main__":
    main()