Problem:


Score : 800 points

Problem Statement

Given is a string S consisting of 0 and 1. Find the number of strings, modulo 998244353, that can result from applying the following operation on S between 0 and K times (inclusive):

Choose a pair of integers i, j (1\leq i < j\leq |S|) such that the i-th and j-th characters of S are 0 and 1, respectively. Remove the j-th character from S and insert it to the immediate left of the i-th character.

Constraints

1 \leq |S| \leq 300

0 \leq K \leq 10^9

S consists of 0 and 1.

Input

Input is given from Standard Input in the following format:

S K

Output

Find the number of strings, modulo 998244353, that can result from applying the operation on S between 0 and K times (inclusive).

Sample Input 1

0101 1

Sample Output 1

4

Four strings, 0101, 0110, 1001, and 1010, can result.

Sample Input 2

01100110 2

Sample Output 2

14

Sample Input 3

1101010010101101110111100011011111011000111101110101010010101010101 20

Sample Output 3

113434815



Sample input:
0101 1


Sample output:
4

Submission status: Accepted

Solution code:
MOD = 998244353


def solve(s, k):
    one_count = []
    ones = 0
    for st in s:
        if st == "1":
            ones += 1
        else:
            one_count.append(ones)
            ones = 0
    one_count.append(ones)
    ss = sum(one_count)
    # print(ss)
    m = len(one_count)
    dp = [[[0] * (ss + 1) for _ in range(ss + 1)] for __ in range(m + 1)]
    dp[0][0][0] = 1
    for i in range(m):
        c = one_count[i]
        for j in range(ss + 1):
            for p in range(ss + 1):
                dp[i][j][p] %= MOD
                for q in range(ss + 1 - max(j, p)):
                    dp[i + 1][j + q][p + q] += dp[i][j][p]
                for q in range(-c, 0):
                    if j + q >= 0:
                        dp[i + 1][j + q][p] += dp[i][j][p]

    # print(dp)
    res = sum([dp[m][0][p] for p in range(min(k + 1, ss + 1))])
    # print(res)
    return res % MOD


def main():
    s, k_ = input().split()
    k = int(k_)
    res = solve(s, k)
    print(res)


def test():
    assert solve("0101", 1) == 4
    assert solve("01100110", 2) == 14
    assert solve("1101010010101101110111100011011111011000111101110101010010101010101", 20) == 113434815


if __name__ == "__main__":
    test()
    main()

Submission status: Accepted

Solution code:
s,k = input().split()
K = int(k)
if s[-1] == "1":
  s += "0"
n = len(s)
ls = []
cnt = 0
for i in range(n):
  if s[i] == "0":
    ls.append(cnt)
    cnt = 0
  else:
    cnt += 1
l = len(ls)
sm = sum(ls)
mod = 998244353
dp = [[[0 for i in range(sm+1)] for j in range(sm+1)] for k in range(l+1)]
dp[0][0][0] = 1
smprv = 0
for i in range(1,l+1):
  num = ls[i-1]
  for j in range(smprv+num,sm+1):
    if j == smprv+num:
      dp[i][j][0] = 1
    for k in range(1,j+1):
      x = 0
      for m in range(j+1):
        if num >= m:
          x += dp[i-1][j-m][k]
        else:
          x += dp[i-1][j-m][k-(m-num)]
      dp[i][j][k] = x%mod
  smprv += num
ans = 0
for i in range(sm+1):
  if i <= K:
    ans += dp[-1][-1][i]
    ans %= mod
print(ans)

Submission status: Accepted

Solution code:
def main():
    mod = 998244353
    s, k = input().split()
    k = int(k)
    n = len(s)

    one = s.count("1")

    cnt = 0
    zero_list = []
    for i in range(n):
        if s[i] == "0":
            zero_list.append(cnt)
            cnt = 0
        else:
            cnt += 1

    z = 0
    mm = min(one, k)

    dp = [[0]*(one+1) for _ in [0]*(one+1)]
    dp[0][0] = 1

    for i in range(len(zero_list)):
        dp2 = [[0]*(mm+1) for _ in [0]*(one+1)]
        base = zero_list[i]
        # j:何個今までに入れたか
        for j in range(one+1):
            # l:何個入れるか
            for l in range(one+1-j):
                if l < z+base-j:
                    continue
                ml = max(l-base, 0)
                # p:これまでのペナルティ
                for p in range(min(one, k)+1):
                    q = p+ml
                    if q <= mm:
                        dp2[j+l][q] = (dp2[j+l][q]+dp[j][p]) % mod
                    else:
                        break
        z += base
        dp = dp2
    print(sum([sum(i) for i in dp]) % mod)


main()

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline

S,K=input().split()
K=int(K)
mod=998244353

LEN=len(S)

ONES=[]
Z=0
for s in S:
    if s=="0":
        Z+=1
    else:
        ONES.append(Z)

if ONES==[]:
    print(1)
    sys.exit()

O=len(ONES)
Z=ONES[-1]

DP=[[0]*(O+1) for i in range(min(O,K)+1)]
DP[0][0]=1

# DP[i][j]で, i個石を動かし、j個石が残っているときの場合の数

from collections import Counter
C=Counter(ONES)

for numb in range(Z,-1,-1):
    #print(DP)
    c=C[numb]

    NDP=[[0]*(O+1) for i in range(min(O,K)+1)]

    for i in range(min(O,K)+1):        
        for j in range(O+1):
            if DP[i][j]==0:
                continue
            for oku in range(c+j+1):
                if i+max(0,oku-c)<=min(O,K) and 0<=c+j-oku<=O:
                    NDP[i+max(0,oku-c)][c+j-oku]=(NDP[i+max(0,oku-c)][c+j-oku]+DP[i][j])%mod
    DP=NDP

ANS=0
for i in range(min(O,K)+1):
    ANS+=DP[i][0]

print(ANS%mod)

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline

S,K=input().split()
K=int(K)
mod=998244353

LEN=len(S)

ONES=[]
Z=0
for s in S:
    if s=="0":
        Z+=1
    else:
        ONES.append(Z)

if ONES==[]:
    print(1)
    sys.exit()

O=len(ONES)
Z=ONES[-1]

DP=[[0]*(O+1) for i in range(min(O,K)+1)]
DP[0][0]=1

# DP[i][j]で, i個石を動かし、j個石が残っているときの場合の数

from collections import Counter
C=Counter(ONES)

for numb in range(Z,-1,-1):
    #print(DP)
    c=C[numb]

    NDP=[[0]*(O+1) for i in range(min(O,K)+1)]

    for i in range(min(O,K)+1):        
        for j in range(O+1):
            if DP[i][j]==0:
                continue
            for oku in range(c+j+1):
                if i+max(0,oku-c)<=min(O,K) and 0<=c+j-oku<=O:
                    NDP[i+max(0,oku-c)][c+j-oku]+=DP[i][j]
    DP=NDP

ANS=0
for i in range(min(O,K)+1):
    ANS+=DP[i][0]

print(ANS%mod)

Submission status: Accepted

Solution code:
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)

Submission status: Accepted

Solution code:
s,n = input().split()
n = int(n)
if n == 0:
    print(1)
    exit(0)
mod = 998244353

b = [0]*301
i = 300
for t in s:
    if t == "0":
        i -= 1
    else:
        b[i] += 1
dp = [[[0]*301 for j in range(301)] for i in range(301)]
sumdelta = [[0]*301 for _ in range(301)]
for j in range(301):
    sumdelta[0][j] = 1
for i in range(301):
    if i == 0:
        dp[i][0][0] = 1
        continue
    delta = [[0]*302 for _ in range(302)]
    for k in range(301):
        inv = [0]*302
        for j in range(k+1)[::-1]:
            inv[j] = inv[j+1] + dp[i-1][j][k]
            inv[j] %= mod
        for j in range(k+1):
            dp[i][j][k] = (inv[j] + sumdelta[k-j][k] - sumdelta[k-j][max(0,k-b[i])])%mod

            delta[k-j][k+1] = (delta[k-j][k]+dp[i][j][k])%mod
    sumdelta = delta.copy()
ans = 0
for i in range(min(n,300)+1):
    ans += dp[300][0][i]
    ans %= mod
print(ans)

Submission status: Accepted

Solution code:
from sys import exit
MOD = 998244353
s, K = input().split()
K = min(int(K), s.count("1"))
if K == 0:
	print(1)
	exit()

a = []
cnt = 0
for i in range(len(s)):
	if s[i] == "0":
		a.append(cnt)
		cnt = 0
	else:
		cnt += 1
a.append(cnt)
a.reverse()
#print(a)
n = len(a)
dp = [[[0 for _ in range(K+1)] for _ in range(K+1)] for _ in range(n+1)]
col = [[0 for _ in range(K+1)] for _ in range(K+1)]
dia = [[0 for _ in range(K+1)] for _ in range(K+1)]
dp[0][0][0] = 1
col[0][0] = 1
for j in range(K+1):
	dia[j][j] = 1

#減らす
def COL(j, k):
	if k == 0:
		return col[j][j]
	else:
		return col[j][j] - col[j][k-1]

#増やす
def DIA(i, j, k):
	if k == 0:
		return 0
	elif k > a[i-1]:
		return dia[j-1][k-1] - dia[j-a[i-1]-1][k-a[i-1]-1]
	else:
		return dia[j-1][k-1]

for i in range(1, n+1):
	for j in range(K+1):
		for k in range(j+1):
			dp[i][j][k] += COL(j, k)
			dp[i][j][k] += DIA(i, j, k)
			dp[i][j][k] %= MOD
	for j in range(K+1):
		col[j][0] = dp[i][j][0]
		dia[j][0] = dp[i][j][0]
		for k in range(1, j+1):
			col[j][k] = (col[j][k-1] + dp[i][j][k]) % MOD
			dia[j][k] = (dia[j-1][k-1] + dp[i][j][k]) % MOD

ans = 0
for j in range(K+1):
	ans += dp[n][j][0]
	ans %= MOD
print(ans)

Submission status: Accepted

Solution code:
from functools import lru_cache
import sys

sys.setrecursionlimit(10 ** 5)
MOD = 998244353


# TLE
def solve(S, K):
    numZeroes = S.count("0")
    # Count number of ones before each zero. This is a unique representation of the string
    ones = [0 for i in range(numZeroes + 1)]
    j = 0
    for x in S:
        if x == "0":
            j += 1
        else:
            assert x == "1"
            ones[j] += 1

    # The ones can move to behind an earlier zero. At most K of them can move.
    # Always best to move directly to desired location to use fewer moves
    @lru_cache(maxsize=None)
    def numWays(i, usedK, extraOnes):
        if usedK == K and extraOnes == 0:
            return 1
        if i == 0:
            # If first zero, have to use up all the extra ones
            return 1

        # Keep current value
        ways = numWays(i - 1, usedK, extraOnes)
        # Use up some extra ones to increment curr
        for j in range(1, extraOnes + 1):
            ways += numWays(i - 1, usedK, extraOnes - j)
            ways %= MOD
        # Steal some ones to decrement curr
        for j in range(1, ones[i] + 1):
            if usedK + j > K:
                break

            ways += numWays(i - 1, usedK + j, extraOnes + j)
            ways %= MOD

        return ways % MOD

    return numWays(len(ones) - 1, 0, 0) % MOD


# TLE
def solve(S, K):
    numZeroes = S.count("0")
    numOnes = len(S) - numZeroes
    # Count number of ones before each zero. This is a unique representation of the string
    ones = [0 for i in range(numZeroes + 1)]
    j = 0
    for x in S:
        if x == "0":
            j += 1
        else:
            assert x == "1"
            ones[j] += 1

    # The ones can move to behind an earlier zero. At most K of them can move.
    # Always best to move directly to desired location to use fewer moves
    K = min(K, numOnes)
    numWays = [
        [[None for extraOnes in range(numOnes + 1)] for usedK in range(K + 1)]
        for i in range(len(ones))
    ]
    for i in range(len(ones)):
        numWays[i][K][0] = 1
    for usedK in range(K + 1):
        for extraOnes in range(numOnes + 1):
            numWays[0][usedK][extraOnes] = 1
    for i in range(len(ones)):
        for usedK in range(K, -1, -1):
            for extraOnes in range(usedK, -1, -1):
                if usedK == K and extraOnes == 0:
                    continue
                if i == 0:
                    continue

                # Keep current value
                ways = numWays[i - 1][usedK][extraOnes]
                # Use up some extra ones to increment curr
                for j in range(1, extraOnes + 1):
                    ways += numWays[i - 1][usedK][extraOnes - j]
                    ways %= MOD
                # Steal some ones to decrement curr
                for j in range(1, ones[i] + 1):
                    if usedK + j > K or extraOnes + j > numOnes:
                        break

                    ways += numWays[i - 1][usedK + j][extraOnes + j]
                    ways %= MOD

                numWays[i][usedK][extraOnes] = ways % MOD

    return numWays[len(ones) - 1][0][0] % MOD


S, K = input().split()
S = list(S)
K = int(K)
print(solve(S, K))

Submission status: Accepted

Solution code:

S,K = input().split()
S = list(S)

ori = [0]
for i in range(len(S)-1,-1,-1):
    if S[i] == "0":
        ori.append(0)
    else:
        ori[-1] += 1

#print (ori)
N = len(ori)

K = int(K)
K = min(310,K)
mod = 998244353

dp = [[[0] * (K+1) for i in range(K+1)] for j in range(N+10)]

dp[-1][0][0] = 1


for i in range(N):

    #なにもしない
    for k in range(K+1):
        for x in range(k+1):
            dp[i][k][x] += dp[i-1][k][x]

    #1を追加する
    for x in range(K+1):

        nsum = 0

        for k in range(K+1):

            if x <= k:
                dp[i][k][x] += nsum
                dp[i][k][x] %= mod

            nsum += dp[i-1][k][x]

            if k-ori[i] >= 0:
                nsum -= dp[i-1][k-ori[i]][x]
            nsum %= mod

    #xを使用する

    for k in range(K+1):

        nsum = 0

        for x in range(k+1):

            dp[i][k][x] += nsum
            dp[i][k][x] %= mod

            if x >= 0:
                nsum += dp[i-1][k][x]

        nsum %= mod

    #print (dp[i])

ans = 0
for i in range(K+1):
    ans += dp[N-1][i][i]
    ans %= mod
print (ans)

Submission status: Accepted

Solution code:
import sys
import numpy as np

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

MOD = 998_244_353

def precompute(S):
    if np.all(S == 1):
        print(1)
        exit()
    while S[0] == 1:
        S = S[1:]
    if np.all(S == 0):
        print(1)
        exit()
    # 各 1 の左側にある 0 の個数を数える
    z = 0
    A = []
    for i, x in enumerate(S):
        if x == 0:
            z += 1
        else:
            A.append(z)
    return np.array(A, dtype=np.int64)

S, K = read().split()
S = np.array(list(S), np.int64) - ord('0')
K = int(K)
A = precompute(S)

A = precompute(S)

def main(A):
    U = 310
    # 数、消費回数、できた禁止地点
    dp = np.zeros((U, U, U), np.int64)
    dp[0, 0, 0] = 1
    for n in range(1, U):
        # 消費回数なし
        dp[n] += dp[n - 1]
        count = np.sum(A == n)
        if not count:
            continue
        tmp = np.zeros_like(dp[n])
        tmp[1:, 1:] += dp[n - 1][:-1, :-1]
        np.cumsum(tmp, axis=1, out=tmp)
        tmp %= MOD
        for k in range(1, count + 1):
            dp[n] += tmp
            tmp2 = np.zeros_like(tmp)
            tmp2[1:] = tmp[:-1]
            np.cumsum(tmp2, axis=1, out=tmp2)
            tmp = tmp2 % MOD
        dp[n, :, n + 1:] = 0
        dp[n] %= MOD
    x = dp[-1][:K + 1, :].sum() % MOD
    return x

print(main(A))

Submission status: Accepted

Solution code:
# coding: utf-8
# Your code here!
import sys
read = sys.stdin.read
readline = sys.stdin.readline

#a,b,c,d = map(int,readline().split())
s,k = readline().split()


a = [len(i) for i in s.split("0")]
while a and a[-1] == 0: a.pop()

if not a:
    print(1)
    exit()

MOD = 998244353

M = sum(a)+1
k = min(int(k),M)


dp = [[0]*M for _ in range(k+1)] # j 使って（上限 k）、l 余ってる
dp[0][0] = 1

#print(a)
for ai in a[::-1]:
    ndp = [[0]*M for _ in range(k+1)] # j 使って（上限 k）、l 余ってる
    for j in range(k+1):
        for l in range(M):
            for ll in range(l):
                ndp[j][ll] += dp[j][l]
                ndp[j][ll] %= MOD
            
            V = min(M-l,k-j+1,ai+1)
            for i in range(V):
                #if j+i > k: break
                ndp[j+i][l+i] += dp[j][l]
                ndp[j+i][l+i] %= MOD
    
    dp = ndp
    #print(dp)

ans = 0
for jj in range(k+1):
    ans += dp[jj][0]

print(ans%MOD)

Submission status: Accepted

Solution code:
P = 998244353
S, K = input().split()
A = [len(a) for a in S.split("0")][::-1]
K = int(K)
M = 301
X = [[0] * M for _ in range(M)]
X[0][0] = 1
for a in A:
    nX = [[0] * M for _ in range(M)]
    Y = [x[:] for x in X]
    for i in range(M-1)[::-1]:
        for j in range(M):
            Y[i][j] += Y[i+1][j]
            if Y[i][j] >= P: Y[i][j] -= P
    Z = [[0] * (M + 1)] + [[0] + x[:] for x in X]
    for i in range(M):
        for j in range(M):
            Z[i+1][j+1] = Z[i][j] + X[i][j]
            if Z[i+1][j+1] >= P: Z[i+1][j+1] -= P
    
    for i in range(M):
        for j in range(M):
            nX[i][j] += Y[i][j]
            t = min(a, i, j)
            nX[i][j] += Z[i][j] - Z[i-t][j-t]
            nX[i][j] %= P
    X = nX
print(sum(X[0][:K+1]) % P)

Submission status: Runtime Error

Solution code:
import sys
readline = sys.stdin.readline
#非再帰
import sys
def scc(Edge):
    N = len(Edge)
    Edgeinv = [[] for _ in range(N)]
    for vn in range(N):
        for vf in Edge[vn]:
            Edgeinv[vf].append(vn)
    
    used = [False]*N
    order = []
    for st in range(N):
        if not used[st]:
            used[st] = True
            stack = [~st, st]
            while stack:
                vn = stack.pop()
                if vn >= 0:
                    for vf in Edge[vn]:
                        if not used[vf]:
                            used[vf] = True
                            stack.append(~vf)
                            stack.append(vf)
                else:
                    order.append(~vn)



    res = [None]*N
    used = [False]*N
    cnt = -1
    for st in order[::-1]:
        if not used[st]:
            cnt += 1
            stack = [st]
            res[st] = cnt
            used[st] = True
            while stack:
                vn = stack.pop()
                for vf in Edgeinv[vn]:
                    if not used[vf]:
                        used[vf] = True
                        res[vf] = cnt
                        stack.append(vf)
    M = cnt+1
    components = [[] for _ in range(M)]
    for i in range(N):
        components[res[i]].append(i)
    """
    tEdge = [[] for _ in range(M)]
    teset = set()
    for vn in range(N):
        tn = res[vn]
        for vf in Edge[vn]:
            tf = res[vf]
            if tn != tf and tn*M + tf not in teset:
                teset.add(tn*M + tf)
                tEdge[tn].append(tf)
    """
    return res, components


        
N, M = map(int, readline().split())
Edge = [[] for _ in range(N)]
for _ in range(M):
    a, b = map(int, readline().split())
    Edge[a].append(b)

R, comp = scc(Edge)
Ans = [str(len(comp))]

for c in comp:
    Ans.append(f'{len(c)} ' + ' '.join(map(str, c)))

print('\n'.join(Ans))

Submission status: Runtime Error

Solution code:
import sys
n=int(input())
alpha=['a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z']
lst=[]
Lst=[]
def Base_10_to_n_(X, n):
    X_dumy = X
    while X_dumy>0:
        Lst.append((X_dumy%n))
        X_dumy = int(X_dumy/n)
    return Lst
if sum(Base_10_to_n_(n,26))==1:
  p=''
  for i in range(len(Base_10_to_n_(n,26))-3):
    p+='z'
  print(p)
  sys.exit()
def Base_10_to_n(X, n):
    X_dumy = X
    ku=0
    while X_dumy>0:
      if ku:
        lst.append(max(0,(X_dumy%n)-1))
        if X_dumy%n==0:
          ku=1
        else:
          ku=0        
        X_dumy = int(X_dumy/n)
      else:
        lst.append((X_dumy%n))
        if X_dumy%n==0:
          ku=1
        else:
          ku=0          
        X_dumy = int(X_dumy/n)
 
    return lst
lst=Base_10_to_n(n,26)
lst.reverse()
t=''
for i in range(len(lst)):
  t+=alpha[lst[i]-1]
print(t)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-

#############
# Libraries #
#############

import sys
input = sys.stdin.readline

import math
#from math import gcd
import bisect
import heapq
from collections import defaultdict
from collections import deque
from collections import Counter
from functools import lru_cache


#############
# Constants #
#############

MOD = 998244353
INF = float('inf')
AZ = "abcdefghijklmnopqrstuvwxyz"

#############
# Functions #
#############

######INPUT######
def I(): return int(input().strip())
def S(): return input().strip()
def IL(): return list(map(int,input().split()))
def SL(): return list(map(str,input().split()))
def ILs(n): return list(int(input()) for _ in range(n))
def SLs(n): return list(input().strip() for _ in range(n))
def ILL(n): return [list(map(int, input().split())) for _ in range(n)]
def SLL(n): return [list(map(str, input().split())) for _ in range(n)]

######OUTPUT######
def P(arg): print(arg); return
def Y(): print("Yes"); return
def N(): print("No"); return
def E(): exit()
def PE(arg): print(arg); exit()
def YE(): print("Yes"); exit()
def NE(): print("No"); exit()

#####Shorten#####
def DD(arg): return defaultdict(arg)


#####Inverse#####
def inv(n): return pow(n, MOD-2, MOD)

######Combination######
kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n):
    return kaijo_memo[n]
  if(len(kaijo_memo) == 0):
    kaijo_memo.append(1)
  while(len(kaijo_memo) <= n):
    kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n):
    return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0):
    gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n):
    gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

def nCr(n,r):
  if(n == r):
    return 1
  if(n < r or r < 0):
    return 0
  ret = 1
  ret = ret * kaijo(n) % MOD
  ret = ret * gyaku_kaijo(r) % MOD
  ret = ret * gyaku_kaijo(n-r) % MOD
  return ret

######Factorization######
def factorization(n):
  arr = []
  temp = n
  for i in range(2, int(-(-n**0.5//1))+1):
    if temp%i==0:
      cnt=0
      while temp%i==0: 
        cnt+=1 
        temp //= i
      arr.append([i, cnt])
  if temp!=1:
    arr.append([temp, 1])
  if arr==[]:
    arr.append([n, 1])
  return arr

#####MakeDivisors######
def make_divisors(n):
  divisors = []
  for i in range(1, int(n**0.5)+1):
    if n % i == 0:
      divisors.append(i)
      if i != n // i: 
        divisors.append(n//i)
  return divisors

#####MakePrimes######
def make_primes(N):
  max = int(math.sqrt(N))
  seachList = [i for i in range(2,N+1)]
  primeNum = []
  while seachList[0] <= max:
    primeNum.append(seachList[0])
    tmp = seachList[0]
    seachList = [i for i in seachList if i % tmp != 0]
  primeNum.extend(seachList)
  return primeNum

#####GCD#####
def gcd(a, b):
    while b:
        a, b = b, a % b
    return a

#####LCM#####
def lcm(a, b):
    return a * b // gcd (a, b)

#####BitCount#####
def count_bit(n):
  count = 0
  while n:
    n &= n -1
    count += 1
  return count

#####ChangeBase#####
def base_10_to_n(X, n):
  if X//n:
    return base_10_to_n(X//n, n)+[X%n]
  return [X%n]

def base_n_to_10(X, n):
  return sum(int(str(X)[-i-1])*n**i for i in range(len(str(X))))

#####IntLog#####
def int_log(n, a):
  count = 0
  while n>=a:
    n //= a
    count += 1
  return count

#############
# Main Code #
#############
S,K = SL()
K = int(K)
LENGTH = 8

data = []

count = 0
for s in S:
  if s == "1":
    count += 1
  else:
    data.append(count)
    count = 0
data.append(count)
data = data[::-1]
N = len(data)
#print(data)

dp = [[[0 for j in range(K+1)] for k in range(K+1)] for i in range(N)]
for i in range(N):
  dp[i][0][0] = 1
for k in range(data[0]+1):
  dp[0][k][k] = 1
#print(dp[0])
for i in range(1,N):
  a = data[i]
  for k in range(1,K+1):   
    for j in range(K+1)[::-1]:
      if j == K:
        dp[i][k][j] = dp[i-1][k][j]
        dp[i][k][j] %= MOD
      else:
        dp[i][k][j] = dp[i-1][k][j] + dp[i][k][j+1]
        dp[i][k][j] %= MOD
  for k in range(K):
    if a:
      q = deque([dp[i-1][k][0]])
      temp = dp[i-1][k][0]
      for j in range(1,K+1):
        #print(i,k,q)
        if k+j > K: break
        dp[i][k+j][j] += temp
        dp[i][k+j][j] %= MOD
        q.append(dp[i-1][k+j][j])
        temp += dp[i-1][k+j][j]
        if len(q) > a:
          temp -= q.popleft()
        temp %= MOD
  #print(dp[i])
      
ans = 0
for k in range(K+1):
  ans += dp[-1][k][0]
  ans %= MOD
  
print(ans)

Submission status: Runtime Error

Solution code:
def main():
    mod = 998244353
    s, k = input().split()
    k, n, one, cnt, z, zero_list = int(k), len(s), s.count("1")+1, 0, 0, []
    for i in range(n):
        if s[i] == "0":
            zero_list.append(cnt)
            cnt = 0
        else:
            cnt += 1
    mm = min(one-1, k)

    dp = [[0]*(mm+1) for _ in [0]*one]
    dp[0][0] = 1

    for i in zero_list:
        dp2 = [[0]*(mm+1) for _ in [0]*one]
        for jl in range(z+i, one):
            for x in range(min(one, jl+1, mm+i+1)):
                dp2[jl][x-i] = (dp2[jl][x-i]+sum([dp[jl-x+p][p]
                                                  for p in range(min(one, x-i+1))])) % mod
            for p in range(min(one, mm+1)):
                dp2[jl][p] = (dp2[jl][p]+sum([dp[jl-l][p]
                                              for l in range(min(one, i, jl-p+1))])) % mod
        z += i
        dp = dp2
    print(sum([sum(i) for i in dp]) % mod)


main()

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-

#############
# Libraries #
#############

import sys
input = sys.stdin.readline

import math
#from math import gcd
import bisect
import heapq
from collections import defaultdict
from collections import deque
from collections import Counter
from functools import lru_cache


#############
# Constants #
#############

MOD = 998244353
INF = float('inf')
AZ = "abcdefghijklmnopqrstuvwxyz"

#############
# Functions #
#############

######INPUT######
def I(): return int(input().strip())
def S(): return input().strip()
def IL(): return list(map(int,input().split()))
def SL(): return list(map(str,input().split()))
def ILs(n): return list(int(input()) for _ in range(n))
def SLs(n): return list(input().strip() for _ in range(n))
def ILL(n): return [list(map(int, input().split())) for _ in range(n)]
def SLL(n): return [list(map(str, input().split())) for _ in range(n)]

######OUTPUT######
def P(arg): print(arg); return
def Y(): print("Yes"); return
def N(): print("No"); return
def E(): exit()
def PE(arg): print(arg); exit()
def YE(): print("Yes"); exit()
def NE(): print("No"); exit()

#####Shorten#####
def DD(arg): return defaultdict(arg)


#####Inverse#####
def inv(n): return pow(n, MOD-2, MOD)

######Combination######
kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n):
    return kaijo_memo[n]
  if(len(kaijo_memo) == 0):
    kaijo_memo.append(1)
  while(len(kaijo_memo) <= n):
    kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n):
    return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0):
    gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n):
    gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

def nCr(n,r):
  if(n == r):
    return 1
  if(n < r or r < 0):
    return 0
  ret = 1
  ret = ret * kaijo(n) % MOD
  ret = ret * gyaku_kaijo(r) % MOD
  ret = ret * gyaku_kaijo(n-r) % MOD
  return ret

######Factorization######
def factorization(n):
  arr = []
  temp = n
  for i in range(2, int(-(-n**0.5//1))+1):
    if temp%i==0:
      cnt=0
      while temp%i==0: 
        cnt+=1 
        temp //= i
      arr.append([i, cnt])
  if temp!=1:
    arr.append([temp, 1])
  if arr==[]:
    arr.append([n, 1])
  return arr

#####MakeDivisors######
def make_divisors(n):
  divisors = []
  for i in range(1, int(n**0.5)+1):
    if n % i == 0:
      divisors.append(i)
      if i != n // i: 
        divisors.append(n//i)
  return divisors

#####MakePrimes######
def make_primes(N):
  max = int(math.sqrt(N))
  seachList = [i for i in range(2,N+1)]
  primeNum = []
  while seachList[0] <= max:
    primeNum.append(seachList[0])
    tmp = seachList[0]
    seachList = [i for i in seachList if i % tmp != 0]
  primeNum.extend(seachList)
  return primeNum

#####GCD#####
def gcd(a, b):
    while b:
        a, b = b, a % b
    return a

#####LCM#####
def lcm(a, b):
    return a * b // gcd (a, b)

#####BitCount#####
def count_bit(n):
  count = 0
  while n:
    n &= n -1
    count += 1
  return count

#####ChangeBase#####
def base_10_to_n(X, n):
  if X//n:
    return base_10_to_n(X//n, n)+[X%n]
  return [X%n]

def base_n_to_10(X, n):
  return sum(int(str(X)[-i-1])*n**i for i in range(len(str(X))))

#####IntLog#####
def int_log(n, a):
  count = 0
  while n>=a:
    n //= a
    count += 1
  return count

#############
# Main Code #
#############
S,K = SL()
K = int(K)
LENGTH = 8

data = []

count = 0
for s in S:
  if s == "1":
    count += 1
  else:
    data.append(count)
    count = 0
data.append(count)
data = data[::-1]
N = len(data)
#print(data)

dp = [[[0 for j in range(K+1)] for k in range(K+1)] for i in range(N)]
for i in range(N):
  dp[i][0][0] = 1
for k in range(data[0]+1):
  dp[0][k][k] = 1
#print(dp[0])
for i in range(1,N):
  a = data[i]
  for k in range(1,K+1):   
    for j in range(K+1):
      for jj in range(K+1):
        if jj < j <= jj+a:
          if k>=(j-jj):
            dp[i][k][j] += dp[i-1][k-(j-jj)][jj]
            dp[i][k][j] %= MOD
        elif j <= jj:
          dp[i][k][j] += dp[i-1][k][jj]
          dp[i][k][j] %= MOD
ans = 0
for k in range(K+1):
  ans += dp[-1][k][0]
  ans %= MOD
  
print(ans)