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: Wrong Answer

Solution code:
from itertools import groupby
import sys
input = sys.stdin.readline

S, K = input().split()
K = min(int(K), S.count('1'))
N = len(S)
mod = 998244353

fact = [1] * (N+1)
fact_inv = [1] * (N+1)
for i in range(1, N+1):
    fact[i] = i * fact[i-1] % mod
fact_inv[N] = pow(fact[N], mod-2, mod)
for i in range(1, N+1)[::-1]:
    fact_inv[i-1] = i * fact_inv[i] % mod

def comb(n, k):
    if k < 0 or n < k:
        return 0
    return fact[n] * fact_inv[k] * fact_inv[n-k] % mod

cnt0 = 0
cnt1 = 0
dp = [[0] * (K+1) for _ in range(K+1)]  # dp[i][j] := moved i 1s from j groups
dp[0][0] = 1
for key, group in groupby(S):
    if key == '0':
        cnt0 += len(list(group))
        continue
    g = len(list(group))
    ndp = [dp[i][:] for i in range(K+1)]
    for i in range(K+1):
        for j in range(K):
            if cnt0 - j <= 0:
                break
            for k in range(1, min(g, K - i) + 1):
                ndp[i+k][j+1] += dp[i][j] * comb(cnt0 - j + k - 1, k)
                ndp[i+k][j+1] %= mod
    cnt1 += 1
    dp = ndp
ans = 0
for i in range(K+1):
    for j in range(K+1):
        ans = (ans + dp[i][j]) % mod
print(ans)

Submission status: Wrong Answer

Solution code:
S,Kstring = input().split()
K = int(Kstring)
lastZero = -1
segment = []
MOD = 998244353
N = len(S)

for i in range(N):
    if S[i] == '0':
        segment.append(i - lastZero - 1)
        lastZero = i

segment.append(N - 1 - lastZero)
M = len(segment)
No1 = N + 1 - M

partialSumSegment = []
for i in range(M):
    if partialSumSegment:
        partialSumSegment.append(segment[i] + partialSumSegment[-1])
    else:
        partialSumSegment.append(segment[i])

DP = []
DPSum = [] # Sum of DP i,j,k ~ i,j,a
for i in range(M):
    DP.append([])
    DPSum.append([])
    for j in range(No1 + 2):
        DP[i].append([])
        DPSum[i].append([])
        for k in range(No1 + 2):
            DP[i][j].append(0)
            DPSum[i][j].append(0)

# Initialize on the last segment

for j in range(No1 - partialSumSegment[M-2] + 1):
    DP[M-1][j][partialSumSegment[M-2] + j] = 1
for j in range(No1 - partialSumSegment[M-2] + 1):
    for k in range(partialSumSegment[M-2] + j,-1,-1):
        DPSum[M-1][j][k] = DPSum[M-1][j][k+1] + DP[M-1][j][k]

for i in range(M-2,-1,-1):
    for j in range(No1 + 1):
        Kmin = partialSumSegment[i-1]
        if i == 0:
            Kmin = 0
        for k in range(No1, Kmin - 1, -1):
            DPCur = 0
            for l in range(1,segment[i] + 1):
                if j >= l and k + segment[i] - l <= No1 + 1:
                    DPCur += DP[i + 1][j - l][k + segment[i] - l]
            if k + segment[i] <= No1:
                DPCur += DPSum[i + 1][j][k + segment[i]]
            DP[i][j][k] = DPCur % MOD
            DPSum[i][j][k] = (DPSum[i][j][k+1] + DP[i][j][k]) % MOD 

# print(DP)
# print(DPSum)

AnsSum = 0
for j in range(min(K,No1) + 1):
    AnsSum += DP[0][j][0]

if M == 0:
    print(1)
else:
    print(AnsSum % MOD)

Submission status: Wrong Answer

Solution code:
S,Kstring = input().split()
K = int(Kstring)
lastZero = -1
segment = []
MOD = 998244353
N = len(S)

for i in range(N):
    if S[i] == '0':
        segment.append(i - lastZero - 1)
        lastZero = i

segment.append(N - 1 - lastZero)
M = len(segment)
No1 = N + 1 - M

partialSumSegment = []
for i in range(M):
    if partialSumSegment:
        partialSumSegment.append(segment[i] + partialSumSegment[-1])
    else:
        partialSumSegment.append(segment[i])

DP = []
DPSum = [] # Sum of DP i,j,k ~ i,j,a
for i in range(M):
    DP.append([])
    DPSum.append([])
    for j in range(No1 + 2):
        DP[i].append([])
        DPSum[i].append([])
        for k in range(No1 + 2):
            DP[i][j].append(0)
            DPSum[i][j].append(0)

# Initialize on the last segment

for j in range(No1 - partialSumSegment[M-2] + 1):
    DP[M-1][j][partialSumSegment[M-2] + j] = 1
for j in range(No1 - partialSumSegment[M-2] + 1):
    for k in range(partialSumSegment[M-2] + j,-1,-1):
        DPSum[M-1][j][k] = DPSum[M-1][j][k+1] + DP[M-1][j][k]

for i in range(M-2,-1,-1):
    for j in range(No1 + 1):
        Kmin = partialSumSegment[i-1]
        if i == 0:
            Kmin = 0
        for k in range(No1, Kmin - 1, -1):
            DPCur = 0
            for l in range(1,segment[i] + 1):
                if j >= l and k + segment[i] - l <= No1:
                    DPCur += DP[i + 1][j - l][k + segment[i] - l]
            if k + segment[i] <= No1:
                DPCur += DPSum[i + 1][j][k + segment[i]]
            DP[i][j][k] = DPCur % MOD
            DPSum[i][j][k] = (DPSum[i][j][k+1] + DP[i][j][k]) % MOD 

# print(DP)
# print(DPSum)

AnsSum = 0
for j in range(min(K,No1) + 1):
    AnsSum += DP[0][j][0]

if M == 0:
    print(1)
else:
    print(AnsSum % MOD)

Submission status: Wrong Answer

Solution code:
S,Kstring = input().split()
K = int(Kstring)
lastZero = -1
segment = []
MOD = 998244353
N = len(S)

for i in range(N):
    if S[i] == '0':
        segment.append(i - lastZero - 1)
        lastZero = i

segment.append(N - 1 - lastZero)
M = len(segment)
No1 = N + 1 - M

partialSumSegment = []
for i in range(M):
    if partialSumSegment:
        partialSumSegment.append(segment[i] + partialSumSegment[-1])
    else:
        partialSumSegment.append(segment[i])

DP = []
DPSum = [] # Sum of DP i,j,k ~ i,j,a
for i in range(M):
    DP.append([])
    DPSum.append([])
    for j in range(No1 + 2):
        DP[i].append([])
        DPSum[i].append([])
        for k in range(No1 + 2):
            DP[i][j].append(0)
            DPSum[i][j].append(0)

# Initialize on the last segment

for j in range(No1 - partialSumSegment[M-2] + 1):
    DP[M-1][j][partialSumSegment[M-2] + j] = 1
for j in range(No1 - partialSumSegment[M-2] + 1):
    for k in range(partialSumSegment[M-2] + j,-1,-1):
        DPSum[M-1][j][k] = DPSum[M-1][j][k+1] + DP[M-1][j][k]

for i in range(M-2,-1,-1):
    for j in range(No1 + 1):
        Kmin = partialSumSegment[i-1]
        if i == 0:
            Kmin = 0
        for k in range(No1, Kmin - 1, -1):
            DPCur = 0
            for l in range(j-segment[i],j):
                if l >= 0 and k + segment[i] - j + l <= No1 and k + segment[i] - j + l  >= 0:
                    DPCur += DP[i + 1][l][k + segment[i] - j + l]
            if k + segment[i] <= No1:
                DPCur += DPSum[i + 1][j][k + segment[i]]
            DP[i][j][k] = DPCur % MOD
            DPSum[i][j][k] = (DPSum[i][j][k+1] + DP[i][j][k]) % MOD 

# print(DP)
# print(DPSum)

AnsSum = 0
for j in range(min(K,No1) + 1):
    AnsSum += DP[0][j][0]

if M == 0:
    print(1)
else:
    print(AnsSum)

Submission status: Wrong Answer

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: Wrong Answer

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: Wrong Answer

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):
            for p in range(ss):
                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: Wrong Answer

Solution code:
import sys,bisect,string,math,time,functools,random,fractions
from heapq import heappush,heappop,heapify
from collections import deque,defaultdict,Counter
from itertools import permutations,combinations,groupby
rep=range
def Golf():n,*t=map(int,open(0).read().split())
def I():return int(input())
def S_():return input()
def IS():return input().split()
def LS():return [i for i in input().split()]
def MI():return map(int,input().split())
def LI():return [int(i) for i in input().split()]
def LI_():return [int(i)-1 for i in input().split()]
def NI(n):return [int(input()) for i in range(n)]
def NI_(n):return [int(input())-1 for i in range(n)]
def StoLI():return [ord(i)-97 for i in input()]
def ItoS(n):return chr(n+97)
def LtoS(ls):return ''.join([chr(i+97) for i in ls])
def RA():return map(int,open(0).read().split())
def GI(V,E,ls=None,Directed=False,index=1):
    org_inp=[];g=[[] for i in range(V)]
    FromStdin=True if ls==None else False
    for i in range(E):
        if FromStdin:
            inp=LI()
            org_inp.append(inp)
        else:
            inp=ls[i]
        if len(inp)==2:
            a,b=inp;c=1
        else:
            a,b,c=inp
        if index==1:a-=1;b-=1
        aa=(a,c);bb=(b,c);g[a].append(bb)
        if not Directed:g[b].append(aa)
    return g,org_inp
def GGI(h,w,search=None,replacement_of_found='.',mp_def={'#':1,'.':0},boundary=1):
    #h,w,g,sg=GGI(h,w,search=['S','G'],replacement_of_found='.',mp_def={'#':1,'.':0},boundary=1) # sample usage
    mp=[boundary]*(w+2);found={}
    for i in range(h):
        s=input()
        for char in search:
            if char in s:
                found[char]=((i+1)*(w+2)+s.index(char)+1)
                mp_def[char]=mp_def[replacement_of_found]
        mp+=[boundary]+[mp_def[j] for j in s]+[boundary]
    mp+=[boundary]*(w+2)
    return h+2,w+2,mp,found
def TI(n):return GI(n,n-1)
def accum(ls):
    rt=[0]
    for i in ls:rt+=[rt[-1]+i]
    return rt
def bit_combination(n,base=2):
    rt=[]
    for tb in range(base**n):s=[tb//(base**bt)%base for bt in range(n)];rt+=[s]
    return rt
def gcd(x,y):
    if y==0:return x
    if x%y==0:return y
    while x%y!=0:x,y=y,x%y
    return y
def YN(x):print(['NO','YES'][x])
def Yn(x):print(['No','Yes'][x])
def show(*inp,end='\n'):
    if show_flg:print(*inp,end=end)

mo=10**9+7
inf=float('inf')
FourNb=[(-1,0),(1,0),(0,1),(0,-1)];EightNb=[(-1,0),(1,0),(0,1),(0,-1),(1,1),(-1,-1),(1,-1),(-1,1)];compas=dict(zip('WENS',FourNb));cursol=dict(zip('LRUD',FourNb))
l_alp=string.ascii_lowercase
#sys.setrecursionlimit(10**7)
read=sys.stdin.buffer.read
readline=sys.stdin.buffer.readline
input=lambda: sys.stdin.readline().rstrip()

class Tree:
    def __init__(self,inp_size=None,ls=None,init=True,index=1):
        self.LCA_init_stat=False
        self.ETtable=[]
        if init:
            if ls==None:
                self.stdin(inp_size,index=index)
            else:
                self.node_size=len(ls)+1
                self.edges,_=GI(self.node_size,self.node_size-1,ls,index=index)
        return
 
    def stdin(self,inp_size=None,index=1):
        if inp_size==None:
            self.node_size=int(input())
        else:
            self.node_size=inp_size
        self.edges,_=GI(self.node_size,self.node_size-1,index=index)
        return
    
    def listin(self,ls,index=0):
        self.node_size=len(ls)+1
        self.edges,_=GI(self.node_size,self.node_size-1,ls,index=index)
        return
 
    def dfs(self,x,func=lambda pr,prv,nx,dist:prv+dist,root_v=0):
        q=deque([x])
        v=[None]*self.node_size
        v[x]=root_v
        while q:
            c=q.pop()
            for nb,d in self.edges[c]:
                if v[nb]==None:
                    q.append(nb)
                    v[nb]=func(c,v[c],nb,d)
        return v
 
    def bfs(self,x,func=lambda pr,prv,nx,dist:prv+dist,root_v=0):
        q=deque([x])
        v=[None]*self.node_size
        v[x]=root_v
        while q:
            c=q.popleft()
            for nb,d in self.edges[c]:
                if v[nb]==None:
                    q.append(nb)
                    v[nb]=func(c,v[c],nb,d)
        return v
 
    def parent(self,x):
        return self.dfs(0,func=lambda pr,prv,nx,dist:pr,root_v=-1)
 
    def topological_sort(self,x):  # return topological sort of the tree
        tps=[]
        q=deque([x])
        v=[None]*self.node_size
        v[x]=0
        while q:
            c=q.popleft()
            tps.append(c)
            for nb,d in self.edges[c]:
                if v[nb]==None:
                    q.append(nb)
                    v[nb]=0
        return tps
 
    def EulerTour(self,x):
        q=deque()
        q.append(x)
        self.depth=[None]*self.node_size
        self.depth[x]=0
        self.ETtable=[]
        self.ETdepth=[]
        self.ETin=[-1]*self.node_size
        self.ETout=[-1]*self.node_size
        cnt=0
        while q:
            c=q.pop()
            if c<0:
                ce=~c
            else:
                ce=c
                for nb,d in self.edges[ce]:
                    if self.depth[nb]==None:
                        q.append(~ce)
                        q.append(nb)
                        self.depth[nb]=self.depth[ce]+1
            self.ETtable.append(ce)
            self.ETdepth.append(self.depth[ce])
            if self.ETin[ce]==-1:
                self.ETin[ce]=cnt
            else:
                self.ETout[ce]=cnt
            cnt+=1
        return
    
    def LCA_init(self,root):
        self.EulerTour(root)
        self.st=SparseTable(self.ETdepth,init_func=min,init_idl=inf)
        #self.st=SegTree(self.node_size*2-1,self.ETdepth,function=min,ide=inf)
        self.LCA_init_stat=True
        return
    
    def LCA(self,root,x,y):
        if self.LCA_init_stat==False:
            self.LCA_init(root)
        xin,xout=self.ETin[x],self.ETout[x]
        yin,yout=self.ETin[y],self.ETout[y]
        a=min(xin,yin)
        b=max(xout,yout,xin,yin)
        id_of_min_dep_in_et=self.st.query_id(a,b+1)
        return self.ETtable[id_of_min_dep_in_et]
 
    def __str__(self):
        return  str(self.edges)
 
    def show(self):
        if all([all([d==1 for nd,d in e]) for e in self.edges]):
            print( [[nd for nd,d in e] for e in self.edges] )
        else:
            print(self)

class SparseTable: # O(N log N) for init, O(1) for query(l,r)
    def __init__(self,ls,init_func=min,init_idl=float('inf')):
        self.func=init_func
        self.idl=init_idl
        self.size=len(ls)
        self.N0=self.size.bit_length()
        self.table=[ls[:]]
        self.index=[list(range(self.size))]
        self.lg=[0]*(self.size+1)
        
        for i in range(2,self.size+1):
            self.lg[i]=self.lg[i>>1]+1
 
        for i in range(self.N0):
            tmp=[self.func(self.table[i][j],self.table[i][min(j+(1<<i),self.size-1)]) for j in range(self.size)]
            tmp_id=[self.index[i][j] if self.table[i][j]==self.func(self.table[i][j],self.table[i][min(j+(1<<i),self.size-1)]) else self.index[i][min(j+(1<<i),self.size-1)] for j in range(self.size)]
            self.table+=[tmp]
            self.index+=[tmp_id]
    
    # return func of [l,r)
    def query(self,l,r):
        if r>self.size:r=self.size
        #N=(r-l).bit_length()-1
        N=self.lg[r-l]
        return self.func(self.table[N][l],self.table[N][max(0,r-(1<<N))])
    
    # return index of which val[i] = func of v among [l,r)
    def query_id(self,l,r):
        if r>self.size:r=self.size
        #N=(r-l).bit_length()-1
        N=self.lg[r-l]
        a,b=self.index[N][l],self.index[N][max(0,r-(1<<N))]
        if self.table[0][a]==self.func(self.table[N][l],self.table[N][max(0,r-(1<<N))]):
            b=a
        return b
    
    # return boundary index of r such that func({table[i]} =< x , i in [pos,r]
    def right_bound(self,pos,x):
        k=(self.size-pos).bit_length()
        for j in range(k)[::-1]:
            nx=pos+(1<<j)
            if nx<n and self.query(pos,nx+1)<=x:
                pos+=(1<<j)
        return pos
    
    # return boundary index of l such that func({table[i]} =< x , i in [l,pos]
    def left_bound(self,pos,x):
        k=pos.bit_length()
        for j in range(k)[::-1]:
            nx=pos-(1<<j)
            if 0<=nx and self.query(nx,pos+1)<=x:
                pos-=(1<<j)
        return pos
 
    def __str__(self):
        return str(self.table[0])
 
    def print(self):
        for i in self.table:
            print(*i)


show_flg=False
show_flg=True


ans=0
mo=998244353

s,k=input().split()
K=int(k)
c=0
t=[]
for i in s+'0':
    if i=='1':
        c+=1
    else:
        t+=[c]
        c=0
t=t[::-1]
one=sum(t)
ans=t
n=len(t)

D=300
m=min(K+1,D)

dp=[[[0]*(m+1) for i in range(one+1)] for i in range(n+1)]
dp[0][0][0]=1

for i in range(n):
    for x in range(m+1):
        for k in range(m+1):
            if 0<=t[i]-x<=one and 0<=k+x<=m:
                print((i+1,t[i]-x,k+x),i,k,x)
                dp[i+1][t[i]-x][k+x]+=dp[i][x][k]

print(ans)

Submission status: Wrong Answer

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)
    s = sum(one_count) // 2 + 1
    m = len(one_count)
    dp = [[[0] * s for _ in range(s)] for __ in range(m + 1)]
    dp[0][0][0] = 1
    for i in range(m):
        c = one_count[i]
        for j in range(s):
            for p in range(s):
                dp[i][j][p] %= MOD
                for q in range(s - 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, s))])
    # 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: Wrong Answer

Solution code:
s,k = input().split()

strings = []

def shift(s, counter):
  if counter == k:
    return
  if s in strings:
    return
  else:
    indexZero = []
    indexOne = []
    for i in range(len(s)):
      if s[i] == "0":
        indexZero.append(i)
      else:
        indexOne.append(i)
    for i in indexZero:
      for j in indexOne:
        if j > i:
          if i > 1:
            foreI =  s[0:i-1]
          elif i == 1:
            foreI = s[0]
          else:
            foreI = ""
          if j < len(s)-1:
            backJ = s[j+1:]
          else:
            backJ = ""
          if i == j-1:
            between = s[i]
          else:
            between = s[i: j]
          newString = foreI + "1" + between + backJ 
          strings.append(newString)
          shift(newString, counter+1)
          
shift(s, 0)
strings.append(s)
output = len(strings)%998244353

print(output)