Problem p02635 (C++)

Buggy submission (status: Runtime Error):
#pragma GCC target("avx2")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("O3")
#include "bits/stdc++.h"
#include <random>
#include <unordered_map>
#include <unordered_set>
using namespace std;
typedef long long ll;
const ll MOD = /*1'000'000'007LL;*/ 998'244'353LL;
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};

string S;
int K;
vector<int> v, sum;

ll dp[305][305][305], memo1[305][305][305], memo2[305][305][305];

ll solve(int, int, int);

ll calc1(int n, int k, int d) {
  if (k < 0 || d < 0)
    return 0;
  if (memo1[n][k][d] != -1)
    return memo1[n][k][d];
  return memo1[n][k][d] = (calc1(n, k - 1, d - 1) + solve(n, k, d)) % MOD;
}

ll calc2(int n, int k, int d) {
  if (k < 0 || d < 0)
    return 0;
  if (memo2[n][k][d] != -1)
    return memo2[n][k][d];
  return memo2[n][k][d] = (calc2(n, k, d - 1) + solve(n, k, d)) % MOD;
}

ll solve(int n, int k, int d) {
  if (k < 0)
    return 0;
  if (n == 0) {
    return 1;
  }
  if (dp[n][k][d] != -1)
    return dp[n][k][d];
  ll ret = 0;
  /*for(int i=0; i<=v[n]; i++){
          ret += solve(n-1, k-v[n]+i, d+i);
  }*/
  ret += calc1(n - 1, k, d + v[n]) - calc1(n - 1, k - v[n] - 1, d - 1) + MOD;
  /*for(int i=v[n]+1; i<=sum[n]-d; i++){
          ret += solve(n-1, k, d+i);
  }*/
  ret += calc2(n - 1, k, sum[n]) - calc2(n - 1, k, d + v[n]) + MOD;
  return dp[n][k][d] = ret % MOD;
}

signed main() {
  cin >> S >> K;
  K = min(K, 302);
  int cnt = 0;
  rep(i, S.size()) {
    if (S[i] == '0') {
      v.pb(cnt);
      cnt = 0;
    } else
      cnt++;
  }
  sum.resize(v.size());
  sum.back() = v.back() + cnt;
  for (int i = v.size() - 2; i >= 0; i--) {
    sum[i] = sum[i + 1] + v[i];
  }

  memset(dp, -1, sizeof(dp));
  memset(memo1, -1, sizeof(memo1));
  memset(memo2, -1, sizeof(memo2));
  ll ans = 0;
  for (int i = 0; i <= cnt; i++) {
    ans += solve(v.size() - 1, K - (cnt - i), i);
  }
  cout << ans % MOD << endl;
}

Fixed submission:
#pragma GCC target("avx2")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("O3")
#include "bits/stdc++.h"
#include <random>
#include <unordered_map>
#include <unordered_set>
using namespace std;
typedef long long ll;
const ll MOD = /*1'000'000'007LL;*/ 998'244'353LL;
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};

string S;
int K;
vector<int> v, sum;

ll dp[305][305][305], memo1[305][305][305], memo2[305][305][305];

ll solve(int, int, int);

ll calc1(int n, int k, int d) {
  if (k < 0 || d < 0)
    return 0;
  if (memo1[n][k][d] != -1)
    return memo1[n][k][d];
  return memo1[n][k][d] = (calc1(n, k - 1, d - 1) + solve(n, k, d)) % MOD;
}

ll calc2(int n, int k, int d) {
  if (k < 0 || d < 0)
    return 0;
  if (memo2[n][k][d] != -1)
    return memo2[n][k][d];
  return memo2[n][k][d] = (calc2(n, k, d - 1) + solve(n, k, d)) % MOD;
}

ll solve(int n, int k, int d) {
  if (k < 0)
    return 0;
  if (n == 0) {
    return 1;
  }
  if (dp[n][k][d] != -1)
    return dp[n][k][d];
  ll ret = 0;
  /*for(int i=0; i<=v[n]; i++){
          ret += solve(n-1, k-v[n]+i, d+i);
  }*/
  ret += calc1(n - 1, k, d + v[n]) - calc1(n - 1, k - v[n] - 1, d - 1) + MOD;
  /*for(int i=v[n]+1; i<=sum[n]-d; i++){
          ret += solve(n-1, k, d+i);
  }*/
  ret += calc2(n - 1, k, sum[n]) - calc2(n - 1, k, d + v[n]) + MOD;
  return dp[n][k][d] = ret % MOD;
}

signed main() {
  cin >> S >> K;
  K = min(K, 302);
  int cnt = 0;
  rep(i, S.size()) {
    if (S[i] == '0') {
      v.pb(cnt);
      cnt = 0;
    } else
      cnt++;
  }
  if (v.size() == 0) {
    cout << 1 << endl;
    return 0;
  }
  sum.resize(v.size());
  sum.back() = v.back() + cnt;
  for (int i = v.size() - 2; i >= 0; i--) {
    sum[i] = sum[i + 1] + v[i];
  }

  memset(dp, -1, sizeof(dp));
  memset(memo1, -1, sizeof(memo1));
  memset(memo2, -1, sizeof(memo2));
  ll ans = 0;
  for (int i = 0; i <= cnt; i++) {
    ans += solve(v.size() - 1, K - (cnt - i), i);
  }
  cout << ans % MOD << endl;
}