Problem p02635 (python)

Fixed submission:
s, k = input().split()
from random import randint
# k = 200
# s = "".join([str(randint(0,1)) for i in range(k)])
# print(s,k)

mod = 998244353
k = min(302, int(k))
n = len(s)
z = s.count("0")
S = [0] * (z+1)
ind = 0
for i in range(n):
    if s[i] == "0":
        ind += 1
    else:
        S[ind] += 1
# print(S)

D = [[0] * (k+1) for i in range(k+1)]
D[0][0] = 1
for ii in range(z+1):
    
    D0 = [[0] * (k+1) for i in range(k+1)]
    for j in range(k):
        if j+1 < k+1:
            for l in range(k+1):
                D0[j+1][l] = (D0[j+1][l] + D[j][l]) % mod
        if S[len(S)-1-ii] + j + 1 < k+1:
            # print("ind",S[len(S)-1-ii] + j + 1)
            for l in range(k+1):
                D0[S[len(S)-1-ii] + j + 1][l] = (D0[S[len(S)-1-ii] + j + 1][l] - D[j][l]) % mod

    for j in range(1, k+1):
        for l in range(k+1):
            D0[j][l] = (D0[j][l] + D0[j-1][l]) % mod

    D1 = [D[i][:] for i in range(k+1)]
    # D1[0][0] = 0
    for i in range(1, k+1):
        r = i-1
        c = 1
        while 0 <= r < k+1 and 0 <= c < k+1:
            # print("r c",r,c)
            D1[r][c] = (D1[r][c] + D1[r+1][c-1]) % mod
            r -= 1
            c += 1
    for i in range(1, k):
        r = k-1
        c = i+1
        while 0 <= r < k+1 and 0 <= c < k+1:
            D1[r][c] = (D1[r][c] + D1[r+1][c-1]) % mod
            r -= 1
            c += 1

    # for i in range(k):
    #     D1[i][0] = 0
    # for i in range(k):
    #     D1[k][i] = 0

    # print(D,D0,D1)
    for i in range(k+1):
        for j in range(k+1):
            D[i][j] = (D0[i][j] + D1[i][j]) % mod

# print(sum(D[i][i] for i in range(k+1)))
print(sum(D[0]) % mod)
# print(D)

Buggy submission (verdict: Runtime Error):
s, k = input().split()
mod = 998244353
k = int(k)
n = len(s)
z = s.count("0")
S = [0] * (z+1)
ind = 0
for i in range(n):
    if s[i] == "0":
        ind += 1
    else:
        S[ind] += 1
# print(S)

D = [[0] * (k+1) for i in range(k+1)]
D[0][0] = 1
for ii in range(z+1):
    
    D0 = [[0] * (k+1) for i in range(k+1)]
    for j in range(k):
        if j+1 < k+1:
            for l in range(k+1):
                D0[j+1][l] = (D0[j+1][l] + D[j][l]) % mod
        if S[len(S)-1-ii] + j + 1 < k+1:
            # print("ind",S[len(S)-1-ii] + j + 1)
            for l in range(k+1):
                D0[S[len(S)-1-ii] + j + 1][l] = (D0[S[len(S)-1-ii] + j + 1][l] - D[j][l]) % mod

    for j in range(1, k+1):
        for l in range(k+1):
            D0[j][l] = (D0[j][l] + D0[j-1][l]) % mod

    D1 = [D[i][:] for i in range(k+1)]
    # D1[0][0] = 0
    for i in range(1, k+1):
        r = i-1
        c = 1
        while 0 <= r < k+1 and 0 <= c < k+1:
            # print("r c",r,c)
            D1[r][c] = (D1[r][c] + D1[r+1][c-1]) % mod
            r -= 1
            c += 1
    for i in range(1, k):
        r = k-1
        c = i+1
        while 0 <= r < k+1 and 0 <= c < k+1:
            D1[r][c] = (D1[r][c] + D1[r+1][c-1]) % mod
            r -= 1
            c += 1

    # for i in range(k):
    #     D1[i][0] = 0
    # for i in range(k):
    #     D1[k][i] = 0

    # print(D,D0,D1)
    for i in range(k+1):
        for j in range(k+1):
            D[i][j] = (D0[i][j] + D1[i][j]) % mod

# print(sum(D[i][i] for i in range(k+1)))
print(sum(D[0]) % mod)
# print(D)