Problem p02635 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep1(i, n) for (int i = 1; i <= (int)(n); i++)

const long long MOD = 998244353;
struct mint {
  long long x;
  mint(long long x = 0) noexcept : x((x + MOD) % MOD) {}
  mint &operator+=(const mint a) {
    if ((x += a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x -= a.x) < 0)
      x += MOD;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= MOD;
    return *this;
  }
  mint operator+(const mint a) const { return mint(*this) += a; }
  mint operator-(const mint a) const { return mint(*this) -= a; }
  mint operator*(const mint a) const { return mint(*this) *= a; }
  mint pow(long long t) const {
    if (t == 0)
      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 { return mint(*this) /= a; }
};

mint dp[305][305][305];

int main() {
  string S;
  cin >> S;
  int K;
  cin >> K;
  vector<int> vec(1);
  int cnt = 0;
  rep(i, S.size()) {
    if (S[i] == '0') {
      cnt++;
      vec.push_back(0);
    } else {
      vec[cnt]++;
    }
  }
  int N = vec.size(), M = S.size() - vec.size() + 1;
  K = min(K, M);
  reverse(vec.begin(), vec.end());
  vector<int> sum_vec(N + 1);
  rep(i, N) sum_vec[i + 1] = sum_vec[i] + vec[i];

  dp[0][0][0] = 1;

  rep(a, N) {
    vector<vector<mint>> sum_a(sum_vec[a + 1] + 1, vector<mint>(K + 1));
    vector<vector<mint>> sum_b(sum_vec[a + 1] + 1, vector<mint>(K + 1));
    rep(b, sum_vec[a + 1] + 1) rep(k, K + 1) {
      sum_a[b][k] = dp[a][b][k];
      if (b > 0)
        sum_a[b][k] += sum_a[b - 1][k];
      if (b >= vec[a])
        sum_a[b][k] -= dp[a][b - vec[a]][k];

      if (b == 0 || k == 0)
        sum_b[b][k] = dp[a][b][k];
      else
        sum_b[b][k] = sum_b[b - 1][k - 1] + dp[a][b][k];
    }

    rep(b, sum_vec[a + 1] + 1) rep(k, K + 1) {
      dp[a + 1][b][k] += sum_a[b][k];
      if (b - vec[a] >= 0)
        dp[a + 1][b][k] += sum_b[b - vec[a]][k];
    }
  }

  mint ans;
  rep(k, K + 1) ans += dp[N][M][k];
  cout << ans.x << endl;
}


Error: 0

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep1(i, n) for (int i = 1; i <= (int)(n); i++)

const long long MOD = 998244353;
long long enc(int a, int b, int c, int d) {
  return (((long long)a * 301 + b) * 301 + c) * 301 + d;
}

vector<long long> dec(long long num) {
  vector<long long> res;
  rep(i, 4) {
    res.push_back(num % 301);
    num /= 301;
  }
  return res;
}

struct mint {
  long long x;
  mint(long long x = 0) noexcept : x((x + MOD) % MOD) {}
  mint &operator+=(const mint a) {
    if ((x += a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x -= a.x) < 0)
      x += MOD;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= MOD;
    return *this;
  }
  mint operator+(const mint a) const { return mint(*this) += a; }
  mint operator-(const mint a) const { return mint(*this) -= a; }
  mint operator*(const mint a) const { return mint(*this) *= a; }
  mint pow(long long t) const {
    if (t == 0)
      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 { return mint(*this) /= a; }
};

int main() {
  string S;
  cin >> S;
  int N = S.size();
  int K;
  cin >> K;
  vector<int> vec(N, 0);
  vector<int> mx;
  int cnt = 0;
  int num1 = 0;
  rep(i, N) {
    if (S[i] == '0')
      cnt++;
    else {
      vec[cnt]++;
      mx.push_back(cnt);
      num1++;
    }
  }
  K = min(K, num1);
  reverse(mx.begin(), mx.end());

  map<long long, mint> mp;
  mp[enc(0, mx[0] + 1, 0, 1)] = 1;

  auto it = mp.begin();
  while (it != mp.end()) {
    long long nown = (*it).first;
    vector<long long> nvec = dec(nown);
    long long n = nvec[3], c = nvec[2], u = nvec[1], k = nvec[0];
    mint now = mp[nown];
    rep(i, c + 1) {
      if (mx[n] < i)
        break;
      if (i == c) {
        if (vec[c] > k)
          mp[enc(n + 1, c, u, k + 1)] += now;
        else if (K > u)
          mp[enc(n + 1, c, u + 1, k)] += now;
      } else {
        if (vec[i] > 0)
          mp[enc(n + 1, i, u, 1)] += now;
        else if (K > u)
          mp[enc(n + 1, i, u + 1, 1)] += now;
      }
    }
    it = mp.find(nown);
    it++;
  }

  mint ans = 0;
  rep(c, N) rep(u, K + 1) rep1(k, N) {
    if (mp.count(enc(num1, c, u, k)) > 0)
      ans += mp[enc(num1, c, u, k)];
  }

  cout << ans.x << endl;
}