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

#include <bits/stdc++.h>
using namespace std;
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(v) begin(v), end(v)
#define fi first
#define se second
template <typename A, typename B> inline bool chmax(A &a, B b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename A, typename B> inline bool chmin(A &a, B b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
using ll = long long;
using pii = pair<int, int>;
constexpr ll INF = 1ll << 30;
constexpr ll longINF = 1ll << 60;
constexpr ll MOD = 998244353;
constexpr bool debug = 0;
//---------------------------------//

template <int M> struct ModInt {
public:
  using value_type = long long;

  ModInt(value_type val = 0) : val(val < 0 ? (M - (-val % M)) % M : val % M) {}

  explicit operator bool() const noexcept { return val; }
  bool operator==(const ModInt &rhs) const noexcept { return val == rhs.val; }
  bool operator!=(const ModInt &rhs) const noexcept { return !(*this == rhs); }
  ModInt operator+() const noexcept { return ModInt(*this); }
  ModInt operator-() const noexcept { return ModInt(0) -= *this; }
  ModInt operator+(const ModInt &rhs) const noexcept {
    return ModInt(*this) += rhs;
  }
  ModInt operator-(const ModInt &rhs) const noexcept {
    return ModInt(*this) -= rhs;
  }
  ModInt operator*(const ModInt &rhs) const noexcept {
    return ModInt(*this) *= rhs;
  }
  ModInt operator/(const ModInt &rhs) const noexcept {
    return ModInt(*this) /= rhs;
  }

  ModInt &operator+=(const ModInt &rhs) noexcept {
    val += rhs.val;
    if (val >= M)
      val -= M;
    return *this;
  }
  ModInt &operator-=(const ModInt &rhs) noexcept {
    if (val < rhs.val)
      val += M;
    val -= rhs.val;
    return *this;
  }
  ModInt &operator*=(const ModInt &rhs) noexcept {
    val = val * rhs.val % M;
    return *this;
  }
  ModInt &operator/=(const ModInt &rhs) noexcept {
    *this *= rhs.inverse();
    return *this;
  }

  ModInt pow(value_type n) const {
    ModInt res = 1, x = val;
    if (n < 0) {
      x = x.inverse();
      n = -n;
    }
    while (n) {
      if (n & 1)
        res *= x;
      x *= x;
      n >>= 1;
    }
    return res;
  }

  ModInt inverse() const {
    long long a = val, a1 = 1, a2 = 0, b = M, b1 = 0, b2 = 1;
    while (b > 0) {
      value_type q = a / b, r = a % b;
      value_type nb1 = a1 - q * b1, nb2 = a2 - q * b2;
      a = b;
      b = r;
      a1 = b1;
      b1 = nb1;
      a2 = b2;
      b2 = nb2;
    }
    assert(a == 1);
    return a1;
  }

  const value_type &get() const noexcept { return val; }
  static decltype(M) get_mod() noexcept { return M; }

  friend std::ostream &operator<<(std::ostream &os, const ModInt &rhs) {
    return os << rhs.val;
  }
  friend std::istream &operator>>(std::istream &is, ModInt &rhs) {
    value_type x;
    is >> x;
    rhs = ModInt(x);
    return is;
  }

private:
  value_type val;
};
using mint = ModInt<MOD>;

int main() {
  string S;
  int K;
  cin >> S >> K;
  chmin(K, S.size());

  vector<int> A;
  A.emplace_back(0);

  for (int i = 0; i < S.size(); ++i) {
    if (S[i] == '1')
      continue;
    int p = 1;
    while (i + p < S.size() && S[i + p] == '1')
      ++p;
    A.emplace_back(p - 1);
    i += p - 1;
  }
  reverse(ALL(A));

  // dp[i][j][k] := i 番目までで合計 j 回 '1' を飛ばして合計 k 回飛ばした '1'
  // を置いた数
  vector<vector<vector<mint>>> dp(
      A.size() + 1,
      vector<vector<mint>>(S.size() + 1, vector<mint>(S.size() + 1)));
  dp[0][0][0] = 1;

  vector<vector<vector<mint>>> sumj(
      A.size() + 1,
      vector<vector<mint>>(S.size() + 2, vector<mint>(S.size() + 2))),
      sumk = sumj;

  sumj[0][1][1] = 1;
  sumk[0][1][1] = 1;
  REP(j, S.size() + 2) REP(k, S.size() + 2) {
    if (j < S.size() + 1)
      sumj[0][j + 1][k] += sumj[0][j][k];
    if (k < S.size() + 1)
      sumk[0][j][k + 1] += sumk[0][j][k];
  }

  FOR(i, 1, A.size() + 1) {
    REP(j, S.size() + 2) {
      REP(k, S.size() + 2) {
        if (j < S.size() + 1)
          sumj[i][j + 1][k] += sumj[i][j][k];
        if (k < S.size() + 1)
          sumk[i][j][k + 1] += sumk[i][j][k];
        if (j == S.size() + 1 || k == S.size() + 1)
          continue;
        if (j < k)
          continue;

        // 前回 c 個 '1' を置いた

        // 前回飛ばした
        // 0 \\leq c \\leq A[i - 1]
        // \\sum dp[i - 1][j - (A[i - 1] - c)][k]
        // = \\sum dp[i - 1][j - 0..A[i-1]][k]

        // 前回飛ばさなかった
        // A[i - 1] < c
        // \\sum dp[i - 1][j][k - (c - A[i - 1])]
        // = \\sum dp[i - 1][j][0..k-1]

        dp[i][j][k] = sumj[i - 1][j + 1][k + 1] -
                      sumj[i - 1][max(j - A[i - 1], 0)][k + 1];
        dp[i][j][k] += sumk[i - 1][j + 1][k];

        sumj[i][j + 1][k + 1] += dp[i][j][k];
        sumk[i][j + 1][k + 1] += dp[i][j][k];

        // printf("%d %d %d %lld\\n", i, j, k, dp[i][j][k].get());
      }
    }
  }

  // for (int u : A) printf("%d ", u);
  // puts("");

  mint ans = 0;
  REP(i, K + 1) ans += dp[A.size()][i][i];
  cout << ans << endl;
}


Error: -11

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(v) begin(v), end(v)
#define fi first
#define se second
template <typename A, typename B> inline bool chmax(A &a, B b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename A, typename B> inline bool chmin(A &a, B b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
using ll = long long;
using pii = pair<int, int>;
constexpr ll INF = 1ll << 30;
constexpr ll longINF = 1ll << 60;
constexpr ll MOD = 998244353;
constexpr bool debug = 0;
//---------------------------------//

template <int M> struct ModInt {
public:
  using value_type = long long;

  ModInt(value_type val = 0) : val(val < 0 ? (M - (-val % M)) % M : val % M) {}

  explicit operator bool() const noexcept { return val; }
  bool operator==(const ModInt &rhs) const noexcept { return val == rhs.val; }
  bool operator!=(const ModInt &rhs) const noexcept { return !(*this == rhs); }
  ModInt operator+() const noexcept { return ModInt(*this); }
  ModInt operator-() const noexcept { return ModInt(0) -= *this; }
  ModInt operator+(const ModInt &rhs) const noexcept {
    return ModInt(*this) += rhs;
  }
  ModInt operator-(const ModInt &rhs) const noexcept {
    return ModInt(*this) -= rhs;
  }
  ModInt operator*(const ModInt &rhs) const noexcept {
    return ModInt(*this) *= rhs;
  }
  ModInt operator/(const ModInt &rhs) const noexcept {
    return ModInt(*this) /= rhs;
  }

  ModInt &operator+=(const ModInt &rhs) noexcept {
    val += rhs.val;
    if (val >= M)
      val -= M;
    return *this;
  }
  ModInt &operator-=(const ModInt &rhs) noexcept {
    if (val < rhs.val)
      val += M;
    val -= rhs.val;
    return *this;
  }
  ModInt &operator*=(const ModInt &rhs) noexcept {
    val = val * rhs.val % M;
    return *this;
  }
  ModInt &operator/=(const ModInt &rhs) noexcept {
    *this *= rhs.inverse();
    return *this;
  }

  ModInt pow(value_type n) const {
    ModInt res = 1, x = val;
    if (n < 0) {
      x = x.inverse();
      n = -n;
    }
    while (n) {
      if (n & 1)
        res *= x;
      x *= x;
      n >>= 1;
    }
    return res;
  }

  ModInt inverse() const {
    long long a = val, a1 = 1, a2 = 0, b = M, b1 = 0, b2 = 1;
    while (b > 0) {
      value_type q = a / b, r = a % b;
      value_type nb1 = a1 - q * b1, nb2 = a2 - q * b2;
      a = b;
      b = r;
      a1 = b1;
      b1 = nb1;
      a2 = b2;
      b2 = nb2;
    }
    assert(a == 1);
    return a1;
  }

  const value_type &get() const noexcept { return val; }
  static decltype(M) get_mod() noexcept { return M; }

  friend std::ostream &operator<<(std::ostream &os, const ModInt &rhs) {
    return os << rhs.val;
  }
  friend std::istream &operator>>(std::istream &is, ModInt &rhs) {
    value_type x;
    is >> x;
    rhs = ModInt(x);
    return is;
  }

private:
  value_type val;
};
using mint = ModInt<MOD>;

int main() {
  string S;
  int K;
  cin >> S >> K;

  chmin(K, S.size());

  reverse(ALL(S));

  vector<int> ppos[2], sum[2];
  REP(i, 2) {
    ppos[i] = vector<int>(S.size(), INF);    // 0-indexed
    sum[i] = vector<int>(S.size() + 1, INF); // 1-indexed
  }

  REP(i, S.size()) ppos[S[i] == '1'][i] = i;
  REP(j, 2)
  for (int i = S.size() - 2; i >= 0; --i) chmin(ppos[j][i], ppos[j][i + 1]);

  REP(i, S.size()) sum[S[i] == '1'][i + 1]++;
  REP(j, 2) REP(i, S.size()) sum[j][i + 1] += sum[j][i];

  vector<vector<vector<mint>>> dp(
      S.size() + 1,
      vector<vector<mint>>(
          S.size() + 1,
          vector<mint>(S.size() + 1))); // [i][j][k] := 長さ i で j 回使える,
                                        // 合計 k 回飛ばしたときのユニーク数

  dp[0][0][0] = 1;
  REP(i, S.size()) REP(j, S.size()) REP(k, K + 1) {
    int p = i + j;
    if (p >= S.size())
      continue;
    if (dp[i][j][k] == 0)
      continue;

    int pos = ppos[0][p];

    if (pos < INF) {
      int plus = sum[1][pos + 1] - sum[1][p] - 1;
      if (k + plus <= K)
        dp[i + 1][j + plus][k] += dp[i][j][k];
    }

    pos = ppos[1][p];
    if (pos < INF) {
      if (pos == p)
        dp[i + 1][j][k] += dp[i][j][k];
      else if (j > 0)
        dp[i + 1][j - 1][k] += dp[i][j][k];
    }
  }

  mint ans = 0;
  REP(i, K + 1) ans += dp[S.size()][0][i];
  cout << ans << endl;
  return 0;
}