Problem p02635 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define rep(i, n) FOR(i, 0, n)
#define rrep(i, n) for (int i = ((int)(n)-1); i >= 0; --i)
#define whole(x) (x).begin(), (x).end()
#define UNIQUE(v) v.erase(unique(v.begin(), v.end()), v.end())
using P = pair<int, int>;
#define debug(var) cerr << "[" << #var << "] " << var << endl
#define chmin(x, y) x = min(x, y)
#define chmax(x, y) x = max(x, y)
const ll mod = 998244353;
const int dx[] = {-1, 0, 1, 0};
const int dy[] = {0, -1, 0, 1};

struct mint {
  ll x; // typedef long long ll;
  mint(ll x = 0) : x((x % mod + mod) % mod) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint a) {
    if ((x += a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += mod - a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const {
    mint res(*this);
    return res += a;
  }
  mint operator-(const mint a) const {
    mint res(*this);
    return res -= a;
  }
  mint operator*(const mint a) const {
    mint res(*this);
    return res *= a;
  }
  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime mod
  mint inv() const { return pow(mod - 2); }
  mint &operator/=(const mint a) { return (*this) *= a.inv(); }
  mint operator/(const mint a) const {
    mint res(*this);
    return res /= a;
  }
};
istream &operator>>(istream &is, const mint &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const mint &a) { return os << a.x; }

struct combination {
  vector<mint> fact, ifact;
  combination(int n) : fact(n + 1), ifact(n + 1) {
    assert(n < mod);
    fact[0] = 1;
    for (int i = 1; i <= n; ++i)
      fact[i] = fact[i - 1] * i;
    ifact[n] = fact[n].inv();
    for (int i = n; i >= 1; --i)
      ifact[i - 1] = ifact[i] * i;
  }
  mint operator()(int n, int k) {
    if (k < 0 || k > n)
      return 0;
    return fact[n] * ifact[k] * ifact[n - k];
  }
};

mint dp[503][503][503];

int main() {
  string s;
  cin >> s;
  int K;
  cin >> K;
  int n = s.size();
  K = min(K, n);
  vector<int> v;
  rep(i, n) {
    if (s[i] == '0')
      v.push_back(i);
  }
  int sv = v.size();
  reverse(whole(v));
  for (int i = 0; i < sv; i++) {
    for (int j = 0; j <= K; j++) {
      for (int k = 0; k <= j; k++) {
        int ini = v[i];
        int now = ini + k;
        if (n - 1 - now < i)
          continue;
        if (now >= n)
          continue;
        // cout << i << " " << j << " " << k << endl;
        if (i == 0) {
          if (j == k)
            dp[i][j][k] = 1;
          // debug(dp[i][j][k]);
          continue;
        }
        int pk = 0;

        chmax(pk, now - v[i - 1] + 1);

        // debug(pk);
        for (int l = pk; l < k; l++) {
          if (j - k + l < 305) {
            dp[i][j][k] += dp[i - 1][j - k + l][l];
          }
        }
        // debug(dp[i][j][k]);
        for (int l = k; l <= K; l++) {
          dp[i][j][k] += dp[i - 1][j][l];
        }
        // debug(dp[i][j][k]);
      }
    }
  }

  mint ans = 0;
  rep(j, 301) rep(k, 301) { ans += dp[sv - 1][j][k]; }

  cout << ans << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define rep(i, n) FOR(i, 0, n)
#define rrep(i, n) for (int i = ((int)(n)-1); i >= 0; --i)
#define whole(x) (x).begin(), (x).end()
#define UNIQUE(v) v.erase(unique(v.begin(), v.end()), v.end())
using P = pair<int, int>;
#define debug(var) cerr << "[" << #var << "] " << var << endl
#define chmin(x, y) x = min(x, y)
#define chmax(x, y) x = max(x, y)
const ll mod = 998244353;
const int dx[] = {-1, 0, 1, 0};
const int dy[] = {0, -1, 0, 1};

struct mint {
  ll x; // typedef long long ll;
  mint(ll x = 0) : x((x % mod + mod) % mod) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint a) {
    if ((x += a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += mod - a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const {
    mint res(*this);
    return res += a;
  }
  mint operator-(const mint a) const {
    mint res(*this);
    return res -= a;
  }
  mint operator*(const mint a) const {
    mint res(*this);
    return res *= a;
  }
  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime mod
  mint inv() const { return pow(mod - 2); }
  mint &operator/=(const mint a) { return (*this) *= a.inv(); }
  mint operator/(const mint a) const {
    mint res(*this);
    return res /= a;
  }
};
istream &operator>>(istream &is, const mint &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const mint &a) { return os << a.x; }

struct combination {
  vector<mint> fact, ifact;
  combination(int n) : fact(n + 1), ifact(n + 1) {
    assert(n < mod);
    fact[0] = 1;
    for (int i = 1; i <= n; ++i)
      fact[i] = fact[i - 1] * i;
    ifact[n] = fact[n].inv();
    for (int i = n; i >= 1; --i)
      ifact[i - 1] = ifact[i] * i;
  }
  mint operator()(int n, int k) {
    if (k < 0 || k > n)
      return 0;
    return fact[n] * ifact[k] * ifact[n - k];
  }
};

mint dp[503][503][503];

int main() {
  string s;
  cin >> s;
  int K;
  cin >> K;
  int n = s.size();
  K = min(K, n);
  vector<int> v;
  rep(i, n) {
    if (s[i] == '0')
      v.push_back(i);
  }
  int sv = v.size();
  reverse(whole(v));
  if (sv == 0 || sv == n) {
    cout << 1 << endl;
    return 0;
  }
  for (int i = 0; i < sv; i++) {
    for (int j = 0; j <= K; j++) {
      for (int k = 0; k <= j; k++) {
        int ini = v[i];
        int now = ini + k;
        if (n - 1 - now < i)
          continue;
        if (now >= n)
          continue;
        // cout << i << " " << j << " " << k << endl;
        if (i == 0) {
          if (j == k)
            dp[i][j][k] = 1;
          // debug(dp[i][j][k]);
          continue;
        }
        int pk = 0;

        chmax(pk, now - v[i - 1] + 1);

        // debug(pk);
        for (int l = pk; l < k; l++) {
          if (j - k + l < 305) {
            dp[i][j][k] += dp[i - 1][j - k + l][l];
          }
        }
        // debug(dp[i][j][k]);
        for (int l = k; l <= K; l++) {
          dp[i][j][k] += dp[i - 1][j][l];
        }
        // debug(dp[i][j][k]);
      }
    }
  }

  mint ans = 0;
  rep(j, 301) rep(k, 301) { ans += dp[sv - 1][j][k]; }

  cout << ans << endl;
  return 0;
}