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: Time Limit Exceeded

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)

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(min(data[0]+1,K+1)):
  dp[0][k][k] = 1
#print(dp[0])
for i in range(1,N):
  for k in range(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
  a = data[i]
  if a:
    for k in range(K):
      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: Time Limit Exceeded

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(i, 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: Time Limit Exceeded

Solution code:
import sys
import math
from collections import defaultdict
from bisect import bisect_left, bisect_right

sys.setrecursionlimit(10**7)
def input():
    return sys.stdin.readline()[:-1]

mod = 998244353

def I(): return int(input())
def LI(): return list(map(int, input().split()))
def LIR(row,col):
    if row <= 0:
        return [[] for _ in range(col)]
    elif col == 1:
        return [I() for _ in range(row)]
    else:
        read_all = [LI() for _ in range(row)]
        return map(list, zip(*read_all))

#################

# nCk
## 0!~n!をmodしつつ求める
def fact_all(n,M=mod):
    f = [1]*(n+1)
    ans = 1
    for i in range(1,n+1):
        ans = (ans*i)%M
        f[i] = ans
    return f

## inv(0!)~inv(n!)をmodしつつ求める
def fact_inv_all(fact_all,M=mod):
    N = len(fact_all)
    finv = [0]*N
    finv[-1] = pow(fact_all[-1],M-2,M)
    for i in range(N-1)[::-1]:
        finv[i] = finv[i+1]*(i+1)%M
    return finv

## nCkをmodしつつ返す
def nCk(n,k,fact_list,inv_list,M=mod):
    return (((fact_list[n]*inv_list[k])%M)*inv_list[n-k])%M

S,K = list(map(str, input().split()))
N = len(S)
K = int(K)

fl = fact_all(301)
il = fact_inv_all(fl)
all_nCk = [[0]*301 for _ in range(301)]
for i in range(301):
    for j in range(301):
        all_nCk[i][j] = nCk(i,j,fl,il)

# 0を境にブロック分け
count = 0
num = []
for i in range(N):
    if S[i] == '1':
        count += 1
    else:
        num.append(count)
        count = 0
num.append(count)

n = len(num)

# dp[i][j][k]: 左からiブロックまで見て，j回操作し，操作禁止がkブロック存在
dp = [[[0]*n for _ in range(301)] for _ in range(n)]

for i in range(n):
    for j in range(301):
        for k in range(n):
            if k >= i+1:
                continue
            elif j == 0:
                if k == 0:
                    dp[i][j][k] = 1
            else:
                dp[i][j][k] = dp[i-1][j][k]
                for l in range(max(0,j-num[i]),j,1):
                    for m in range(k):
                        dp[i][j][k] += dp[i-1][l][m]*all_nCk[j+k-(l+m+2)][j-(l+1)]
                        dp[i][j][k] %= mod

ans = 0
for j in range(min(301,K+1)):
    for k in range(n):
        ans += dp[n-1][j][k]
        ans %= mod

print(ans)

Submission status: Time Limit Exceeded

Solution code:
import sys
import math
from collections import defaultdict
from bisect import bisect_left, bisect_right

sys.setrecursionlimit(10**7)
def input():
    return sys.stdin.readline()[:-1]

mod = 998244353

def I(): return int(input())
def LI(): return list(map(int, input().split()))
def LIR(row,col):
    if row <= 0:
        return [[] for _ in range(col)]
    elif col == 1:
        return [I() for _ in range(row)]
    else:
        read_all = [LI() for _ in range(row)]
        return map(list, zip(*read_all))

#################

# nCk
## 0!~n!をmodしつつ求める
def fact_all(n,M=mod):
    f = [1]*(n+1)
    ans = 1
    for i in range(1,n+1):
        ans = (ans*i)%M
        f[i] = ans
    return f

## inv(0!)~inv(n!)をmodしつつ求める
def fact_inv_all(fact_all,M=mod):
    N = len(fact_all)
    finv = [0]*N
    finv[-1] = pow(fact_all[-1],M-2,M)
    for i in range(N-1)[::-1]:
        finv[i] = finv[i+1]*(i+1)%M
    return finv

## nCkをmodしつつ返す
def nCk(n,k,fact_list,inv_list,M=mod):
    return (((fact_list[n]*inv_list[k])%M)*inv_list[n-k])%M

S,K = list(map(str, input().split()))
N = len(S)
K = int(K)

fl = fact_all(305)
il = fact_inv_all(fl)
all_nCk = [[0]*305 for _ in range(305)]
for i in range(305):
    for j in range(305):
        all_nCk[i][j] = nCk(i,j,fl,il)

# 0を境にブロック分け
count = 0
num = []
for i in range(N):
    if S[i] == '1':
        count += 1
    else:
        num.append(count)
        count = 0
num.append(count)

n = len(num)

# dp[i][j][k]: 左からiブロックまで見て，j回操作し，操作禁止がkブロック存在
dp = [[[0]*n for _ in range(305)] for _ in range(n)]

for i in range(n):
    for j in range(305):
        for k in range(n):
            if k >= i+1:
                continue
            elif j == 0:
                if k == 0:
                    dp[i][j][k] = 1
            else:
                dp[i][j][k] = dp[i-1][j][k]
                for l in range(max(0,j-num[i]),j,1):
                    for m in range(k):
                        dp[i][j][k] += dp[i-1][l][m]*all_nCk[j+k-(l+m+2)][j-(l+1)]
                        dp[i][j][k] %= mod

ans = 0
for j in range(min(305,K+1)):
    for k in range(n):
        ans += dp[n-1][j][k]
        ans %= mod

print(ans)

Submission status: Time Limit Exceeded

Solution code:
import sys
import math
from collections import defaultdict
from bisect import bisect_left, bisect_right

sys.setrecursionlimit(10**7)
def input():
    return sys.stdin.readline()[:-1]

mod = 998244353

def I(): return int(input())
def LI(): return list(map(int, input().split()))
def LIR(row,col):
    if row <= 0:
        return [[] for _ in range(col)]
    elif col == 1:
        return [I() for _ in range(row)]
    else:
        read_all = [LI() for _ in range(row)]
        return map(list, zip(*read_all))

#################

# nCk
## 0!~n!をmodしつつ求める
def fact_all(n,M=mod):
    f = [1]*(n+1)
    ans = 1
    for i in range(1,n+1):
        ans = (ans*i)%M
        f[i] = ans
    return f

## inv(0!)~inv(n!)をmodしつつ求める
def fact_inv_all(fact_all,M=mod):
    N = len(fact_all)
    finv = [0]*N
    finv[-1] = pow(fact_all[-1],M-2,M)
    for i in range(N-1)[::-1]:
        finv[i] = finv[i+1]*(i+1)%M
    return finv

## nCkをmodしつつ返す
def nCk(n,k,fact_list,inv_list,M=mod):
    return (((fact_list[n]*inv_list[k])%M)*inv_list[n-k])%M

S,K = list(map(str, input().split()))
N = len(S)
K = int(K)

fl = fact_all(305)
il = fact_inv_all(fl)

# 0を境にブロック分け
count = 0
num = []
for i in range(N):
    if S[i] == '1':
        count += 1
    else:
        num.append(count)
        count = 0
num.append(count)

n = len(num)

# dp[i][j][k]: 左からiブロックまで見て，j回操作し，操作禁止がkブロック存在
dp = [[[0]*n for _ in range(305)] for _ in range(n)]

for i in range(n):
    for j in range(305):
        for k in range(n):
            if k >= i+1:
                continue
            elif j == 0:
                if k == 0:
                    dp[i][j][k] = 1
            else:
                dp[i][j][k] = dp[i-1][j][k]
                for l in range(max(0,j-num[i]),j,1):
                    for m in range(k):
                        dp[i][j][k] += dp[i-1][l][m]*nCk(j+k-(l+m+2),j-(l+1),fl,il)
                        dp[i][j][k] %= mod

ans = 0
for j in range(min(305,K+1)):
    for k in range(n):
        ans += dp[n-1][j][k]
        ans %= mod

print(ans)

Submission status: Time Limit Exceeded

Solution code:
mod = 998244353
S,K = input().split()
K = min(300,int(K))
a = []
c = 0
for i in range(len(S)):
    if S[i]=='0':
        a.append(c)
        c = 0
    else:
        c += 1
if S[-1]=='1':
    a.append(c)
a = a[::-1]
N = len(a)
dp = [[[0]*(K+1) for _ in range(K+1)] for _ in range(N+1)]
dp[0][0][0] = 1
for i in range(N):
    for j in range(K+1):
        for k in range(j+1):
            for l in range(a[i]+1):
                if j+l>K:
                    break
                dp[i+1][j+l][k+l] += dp[i][j][k]
                dp[i+1][j+l][k+l] %= mod
            for l in range(1,k+1):
                dp[i+1][j][k-l] += dp[i][j][k]
                dp[i+1][j][k-l] %= mod
                
ans = 0
for k in range(K+1):
    ans += dp[N][k][0]
    ans %= mod
print(ans)

Submission status: Time Limit Exceeded

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)
    ans = 0

    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 j in range(one):
            for l in range(max(z+i-j, 0), one-j):
                ml = max(l-i, 0)
                for p in range(min(j, mm-ml)+1):
                    ans += 1
                    dp2[j+l][p+ml] = (dp2[j+l][p+ml]+dp[j][p]) % mod
        z += i
        dp = dp2
    print(sum([sum(i) for i in dp]) % mod)


main()

Submission status: Time Limit Exceeded

Solution code:
S,K=input().split()
K=int(K)
K=min(K,300)
mod=998244353
a=[]
val=0
for i in range(len(S)):
    if S[i]=="0":
        a.append(val)
        val=0
    else:
        val+=1

if val!=0:
    a.append(val)
m=len(a)

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

for j in range(K+1):
    dp[j][j]=1

for i in range(m-1,-1,-1):
    ndp=[[0 for x in range(K+1)] for y in range(K+1)]
    for j in range(K,-1,-1):
        for k in range(K,-1,-1):
            M=max(k-j,-a[i])
            for l in range(max(0,M),K-j+1):
                ndp[j][k]+=dp[j+l][k]
                ndp[j][k]%=mod
            for l in range(1,min(K-k,-M)+1):
                ndp[j][k]+=dp[j][k+l]
                ndp[j][k]%=mod
    dp=ndp




print(dp[0][0])

Submission status: Time Limit Exceeded

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):
      for m in range(j+1):
        if num >= m:
          dp[i][j][k] += dp[i-1][j-m][k]
        else:
          dp[i][j][k] += dp[i-1][j-m][k-(m-num)]
        dp[i][j][k] %= 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: Time Limit Exceeded

Solution code:
s,k = input().split()
mod = 998244353
k = int(k)
k = min(200,k)
raw = [0]
l = len(s)
for i in range(l):
  if s[i] == '0':
    raw.append(0)
  else:
    raw[-1] += 1
#print(raw)
now = [[0 for _ in range(k+1)] for _ in range(k+1)]
l = len(raw)
raw.reverse()
#print(raw)
now[0][0] = 1
for x in raw:
  last = now[:]
  now = [[0 for _ in range(k+1)] for _ in range(k+1)]
  for i in range(k+1):
    for j in range(k+1):
      for mag in range(x+1):
        if mag + i > k or mag + j > k:
          continue
        now[i+mag][j+mag] += last[i][j]
        now[i+mag][j+mag] %= mod
      for dim in range(j):
        now[i][dim] += last[i][j]
        now[i][dim] %= mod
  #print(now)
#print(now)
ans = 0
for i in range(k+1):
  ans += now[i][0]
print(ans%mod)

Submission status: Time Limit Exceeded

Solution code:
s,k = input().split()
mod = 998244353
k = int(k)
k = min(300,k)
raw = [0]
l = len(s)
for i in range(l):
  if s[i] == '0':
    raw.append(0)
  else:
    raw[-1] += 1
#print(raw)
now = [[0 for _ in range(k+1)] for _ in range(k+1)]
l = len(raw)
raw.reverse()
#print(raw)
now[0][0] = 1
for x in raw:
  last = now[:]
  now = [[0 for _ in range(k+1)] for _ in range(k+1)]
  for i in range(k+1):
    for j in range(k+1):
      for mag in range(x+1):
        if mag + i > k or mag + j > k:
          continue
        now[i+mag][j+mag] += last[i][j]
        now[i+mag][j+mag] %= mod
      for dim in range(j):
        now[i][dim] += last[i][j]
        now[i][dim] %= mod
  #print(now)
#print(now)
ans = 0
for i in range(k+1):
  ans += now[i][0]
print(ans%mod)

Submission status: Time Limit Exceeded

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):
      for m in range(j+1):
        if num >= m:
          dp[i][j][k] += dp[i-1][j-m][k]
        else:
          dp[i][j][k] += dp[i-1][j-m][k-(m-num)]
      dp[i][j][k] %= 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: Time Limit Exceeded

Solution code:
S,K=input().split()
K=int(K)
K=min(K,300)
mod=998244353
a=[]
val=0
for i in range(len(S)):
    if S[i]=="0":
        a.append(val)
        val=0
    else:
        val+=1

if val!=0:
    a.append(val)
m=len(a)

dp=[[[0 for i in range(K+1)] for j in range(K+1)] for k in range(m+1)]

for j in range(K+1):
    dp[m][j][j]=1

for i in range(m-1,-1,-1):
    for j in range(K+1):
        for k in range(K+1):
            M=max(k-j,-a[i])
            for l in range(max(0,M),K-j+1):
                dp[i][j][k]+=dp[i+1][j+l][k]
                dp[i][j][k]%=mod
            for l in range(1,min(K-k,-M)+1):
                dp[i][j][k]+=dp[i+1][j][k+l]
                dp[i][j][k]%=mod




print(dp[0][0][0])

Submission status: Time Limit Exceeded

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

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

    for i in range(len(zero_list)):
        dp2 = [[0]*(min(one, k)+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 j+l > one:
                    continue
                if j+l < z+base:
                    continue
                # p:これまでのペナルティ
                for p in range(min(one, k)+1):
                    q = p+max(l-base, 0)
                    if q <= min(one, k) and j+l <= one:
                        dp2[j+l][q] = (dp2[j+l][q]+dp[j][p]) % mod
        z += base
        dp = dp2
    print(sum([sum(i) for i in dp]) % mod)


main()

Submission status: Time Limit Exceeded

Solution code:

def main():
    S,K=input().split()
    K=int(K)
    K=min(K,300)
    mod=998244353
    a=[]
    val=0
    for i in range(len(S)):
        if S[i]=="0":
            a.append(val)
            val=0
        else:
            val+=1
    
    if val!=0:
        a.append(val)
    m=len(a)
    
    dp=[[[0 for i in range(K+1)] for j in range(K+1)] for k in range(m+1)]
    
    for j in range(K+1):
        dp[m][j][j]=1
    
    for i in range(m-1,-1,-1):
        for j in range(K+1):
            for k in range(K+1):
                M=max(k-j,-a[i])
                if M>=0:
                    for l in range(max(0,M),K-j+1):
                        dp[i][j][k]+=dp[i+1][j+l][k]
                        dp[i][j][k]%=mod
                else:
                    for l in range(max(0,M),K-j+1):
                        dp[i][j][k]+=dp[i+1][j+l][k]
                        dp[i][j][k]%=mod
                    for l in range(1,min(K-k,-M)+1):
                        dp[i][j][k]+=dp[i+1][j][k+l]
                        dp[i][j][k]%=mod
    
    print(dp[0][0][0])

if __name__=="__main__":
    main()

Submission status: Time Limit Exceeded

Solution code:
S,K=input().split()
K=int(K)
K=min(K,300)
mod=998244353
a=[]
val=0
for i in range(len(S)):
    if S[i]=="0":
        a.append(val)
        val=0
    else:
        val+=1

if val!=0:
    a.append(val)
m=len(a)

dp=[[[0 for i in range(K+1)] for j in range(K+1)] for k in range(m+1)]

for j in range(K+1):
    dp[m][j][j]=1

for i in range(m-1,-1,-1):
    for j in range(K+1):
        for k in range(K+1):
            M=max(k-j,-a[i])
            if M>=0:
                for l in range(max(0,M),K-j+1):
                    dp[i][j][k]+=dp[i+1][j+l][k]
                    dp[i][j][k]%=mod
            else:
                for l in range(max(0,M),K-j+1):
                    dp[i][j][k]+=dp[i+1][j+l][k]
                    dp[i][j][k]%=mod
                for l in range(1,min(K-k,-M)+1):
                    dp[i][j][k]+=dp[i+1][j][k+l]
                    dp[i][j][k]%=mod

print(dp[0][0][0])

Submission status: Time Limit Exceeded

Solution code:
S,K=input().split()
K=int(K)
K=min(K,300)
mod=998244353
a=[]
val=0
for i in range(len(S)):
    if S[i]=="0":
        a.append(val)
        val=0
    else:
        val+=1

if val!=0:
    a.append(val)
m=len(a)

dp=[[[0 for i in range(K+1)] for j in range(K+1)] for k in range(m+1)]

for j in range(K+1):
    dp[m][j][j]=1

for i in range(m-1,-1,-1):
    for j in range(K+1):
        for k in range(K+1):
            M=max(k-j,-a[i])
            if M>=0:
                dp[i][j][k]=sum(dp[i+1][j+l][k] for l in range(max(0,M),K-j+1))
            else:
                dp[i][j][k]=sum(dp[i+1][j+l][k] for l in range(max(0,M),K-j+1))+sum(dp[i+1][j][k+l] for l in range(1,min(K-k,-M)+1))
            dp[i][j][k]%=mod

print(dp[0][0][0])

Submission status: Time Limit Exceeded

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
    @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


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

    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(numOnes, -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: Time Limit Exceeded

Solution code:
from functools import lru_cache
import sys

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


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
    @lru_cache(maxsize=None)
    def numWays(i, k, extraOnes):
        if k == 0 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, k, extraOnes)
        # Use up some extra ones to increment curr
        for j in range(1, extraOnes + 1):
            ways += numWays(i - 1, k, extraOnes - j)
            ways %= MOD
        # Steal some ones to decrement curr
        for j in range(1, k + 1):
            if ones[i] - j < 0:
                break
            ways += numWays(i - 1, k - j, extraOnes + j)
            ways %= MOD

        return ways % MOD

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


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

Submission status: Time Limit Exceeded

Solution code:
S, K = input().split()
K = int(K)


def get_variants(variants):
    new_variants = set()
    for variant in variants:
        for i in range(len(variant) - 1):
            if variant[i] == "0":
                for j in range(i + 1, len(variant)):
                    if variant[j] == "1":
                        new_variant = variant[:i] + "1" + variant[i:j] + variant[j+1:]
                        if new_variant not in variants:
                            new_variants.add(new_variant)
    return new_variants


def reorder(s, k):
    variants = {s}
    new_variants = {s}
    for i in range(k):
        new_variants = get_variants(new_variants)
        variants |= new_variants

    return len(variants) % 998244353


print(reorder(S, K))

Submission status: Accepted

Solution code:
import numpy as np


def solve(slots_array, K):
    MOD = 998244353
    dp = np.zeros((K+1, K+1), np.int64)
    dp[0][0] = 1
    for i in range(len(slots_array) - 1, -1, -1):
        new_dp = dp.copy()
        for j in range(K, -1, -1):
            for k in range(K-1, -1, -1):
                new_dp[j][k] += new_dp[j][k+1]
        new_dp %= MOD
        for t in range(1, slots_array[i] + 1):
            new_dp[t:, t:] += dp[:-t, :-t]
        dp = new_dp
    # print(dp)
    return np.sum(dp[:K+1, 0]) % MOD

def main():     
    S, K_str = input().split()
    K = int(K_str)

    slots = []
    index = 0
    while index < len(S):
        prev_index = index
        while index < len(S) and S[index] == '1':
            index += 1
        slots.append(index - prev_index)
        index += 1

    K = min(sum(slots) + 1, K)
    print(solve(np.array(slots, np.int64), K))


def cc_export():
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('solve', 'i8(i8[:], i8)')(solve)
    cc.compile()

if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
import numpy as np


def solve(slots_array, K):
    MOD = 998244353
    dp = np.zeros((K+1, K+1), np.int64)
    dp[0][0] = 1
    for i in range(len(slots_array) - 1, -1, -1):
        new_dp = dp.copy()
        for j in range(K, -1, -1):
            for k in range(K-1, -1, -1):
                new_dp[j][k] += new_dp[j][k+1]
        new_dp %= MOD
        for t in range(1, slots_array[i] + 1):
            new_dp[t:, t:] += dp[:-t, :-t]
        dp = new_dp
    # print(dp)
    return np.sum(dp[:K+1, 0]) % MOD

def main():     
    S, K_str = input().split()
    K = int(K_str)

    slots = []
    index = 0
    while index < len(S):
        prev_index = index
        while index < len(S) and S[index] == '1':
            index += 1
        slots.append(index - prev_index)
        index += 1

    K = min(sum(slots) + 1, K)
    print(solve(np.array(slots, np.int64), K))


def cc_export():
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('solve', 'i8(i8[:], i8)')(solve)
    cc.compile()

if __name__ == '__main__':
    import sys
    if sys.argv[-1] == 'ONLINE_JUDGE':
        cc_export()
        exit(0)
    from my_module import solve
    main()