Problem p02635 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, b) for (int i = (a); i < (b); ++i)
#define per(i, a, b) for (int i = (b)-1; i >= (a); --i)
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define fst first
#define snd second

template <class T> bool ckmin(T &a, const T &b) { return b < a ? a = b, 1 : 0; }
template <class T> bool ckmax(T &a, const T &b) { return a < b ? a = b, 1 : 0; }

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<pii> vii;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

ll dp[302][302][302];

int main() {
  cin.tie(0)->sync_with_stdio(0);
  cin.exceptions(cin.failbit);

  memset(dp, 0ll, sizeof dp);

  int K;
  string s;
  cin >> s >> K;

  s.pb('0');

  vi a;
  int cnt = 0;
  for (auto x : s) {
    if (x == '0') {
      a.pb(cnt);
      cnt = 0;
    } else {
      ++cnt;
    }
  }

  int n = sz(a);
  vi pref = a;
  rep(i, 1, n) pref[i] += pref[i - 1];

  const ll mod = 998244353;

  // dp[i][j][k] => # sequences w [0...i) having prefix sum j and k needed
  // operations
  dp[0][0][0] = 1;
  rep(i, 0, n) {
    rep(j, 0, 301) {
      rep(k, 0, K + 1) {
        if (dp[i][j][k] == 0)
          continue;
        rep(nxt, max(0, pref[i] - j), pref[n - 1] - j + 1) {
          if (max(0, nxt - a[i]) + k > K)
            break;
          dp[i + 1][j + nxt][k + max(0, nxt - a[i])] += dp[i][j][k];
          dp[i + 1][j + nxt][k + max(0, nxt - a[i])] %= mod;
        }
      }
    }
  }

  ll ans = 0;
  rep(j, 0, 301) {
    rep(k, 0, K + 1) {
      ans += dp[n][j][k];
      ans %= mod;
    }
  }
  cout << ans << '\n';
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, b) for (int i = (a); i < (b); ++i)
#define per(i, a, b) for (int i = (b)-1; i >= (a); --i)
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define fst first
#define snd second

template <class T> bool ckmin(T &a, const T &b) { return b < a ? a = b, 1 : 0; }
template <class T> bool ckmax(T &a, const T &b) { return a < b ? a = b, 1 : 0; }

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<pii> vii;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

ll dp[302][302][302];

int main() {
  cin.tie(0)->sync_with_stdio(0);
  cin.exceptions(cin.failbit);

  memset(dp, 0ll, sizeof dp);

  int K;
  string s;
  cin >> s >> K;

  ckmin(K, 300);
  s.pb('0');

  vi a;
  int cnt = 0;
  for (auto x : s) {
    if (x == '0') {
      a.pb(cnt);
      cnt = 0;
    } else {
      ++cnt;
    }
  }

  int n = sz(a);
  vi pref = a;
  rep(i, 1, n) pref[i] += pref[i - 1];

  const ll mod = 998244353;

  // dp[i][j][k] => # sequences w [0...i) having prefix sum j and k needed
  // operations
  dp[0][0][0] = 1;
  rep(i, 0, n) {
    rep(j, 0, 301) {
      rep(k, 0, K + 1) {
        if (dp[i][j][k] == 0)
          continue;
        rep(nxt, max(0, pref[i] - j), pref[n - 1] - j + 1) {
          if (max(0, nxt - a[i]) + k > K)
            break;
          dp[i + 1][j + nxt][k + max(0, nxt - a[i])] += dp[i][j][k];
          dp[i + 1][j + nxt][k + max(0, nxt - a[i])] %= mod;
        }
      }
    }
  }

  ll ans = 0;
  rep(j, 0, 301) {
    rep(k, 0, K + 1) {
      ans += dp[n][j][k];
      ans %= mod;
    }
  }
  cout << ans << '\n';
}