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

#include <bits/stdc++.h>

int ri() {
  int n;
  scanf("%d", &n);
  return n;
}

template <int mod> struct ModInt {
  int x;
  ModInt() : x(0) {}
  ModInt(int64_t x) : x(x >= 0 ? x % mod : (mod - -x % mod) % mod) {}
  ModInt &operator+=(const ModInt &p) {
    if ((x += p.x) >= mod)
      x -= mod;
    return *this;
  }
  ModInt &operator-=(const ModInt &p) {
    if ((x += mod - p.x) >= mod)
      x -= mod;
    return *this;
  }
  ModInt &operator*=(const ModInt &p) {
    x = (int64_t)x * p.x % mod;
    return *this;
  }
  ModInt &operator/=(const ModInt &p) {
    *this *= p.inverse();
    return *this;
  }
  ModInt &operator^=(int64_t p) {
    ModInt res = 1;
    for (; p; p >>= 1) {
      if (p & 1)
        res *= *this;
      *this *= *this;
    }
    return *this = res;
  }
  ModInt operator-() const { return ModInt(-x); }
  ModInt operator+(const ModInt &p) const { return ModInt(*this) += p; }
  ModInt operator-(const ModInt &p) const { return ModInt(*this) -= p; }
  ModInt operator*(const ModInt &p) const { return ModInt(*this) *= p; }
  ModInt operator/(const ModInt &p) const { return ModInt(*this) /= p; }
  ModInt operator^(int64_t p) const { return ModInt(*this) ^= p; }
  bool operator==(const ModInt &p) const { return x == p.x; }
  bool operator!=(const ModInt &p) const { return x != p.x; }
  explicit operator int() const { return x; }
  ModInt &operator=(const int p) {
    x = p;
    return *this;
  }
  ModInt inverse() const {
    int a = x, b = mod, u = 1, v = 0, t;
    while (b > 0) {
      t = a / b;
      a -= t * b;
      std::swap(a, b);
      u -= t * v;
      std::swap(u, v);
    }
    return ModInt(u);
  }
  friend std::ostream &operator<<(std::ostream &stream, const ModInt<mod> &p) {
    return stream << p.x;
  }
  friend std::istream &operator>>(std::istream &stream, ModInt<mod> &a) {
    int64_t x;
    stream >> x;
    a = ModInt<mod>(x);
    return stream;
  }
};

template <int mod> struct MComb {
  using mint = ModInt<mod>;
  std::vector<mint> fact;
  std::vector<mint> inv;
  MComb(int n) { // O(n + log(mod))
    fact = std::vector<mint>(n + 1, 1);
    for (int i = 1; i <= n; i++)
      fact[i] = fact[i - 1] * mint(i);
    inv.resize(n + 1);
    inv[n] = fact[n] ^ (mod - 2);
    for (int i = n; i--;)
      inv[i] = inv[i + 1] * mint(i + 1);
  }
  mint ncr(int n, int r) { return fact[n] * inv[r] * inv[n - r]; }
  mint npr(int n, int r) { return fact[n] * inv[n - r]; }
  mint nhr(int n, int r) {
    assert(n + r - 1 < (int)fact.size());
    return ncr(n + r - 1, r);
  }
};

typedef ModInt<998244353> mint;

int main() {
  std::string s;
  std::cin >> s;
  int n = s.size();

  std::vector<std::vector<int>> ok(
      n + 1, std::vector<int>(n + 1, -1)); // white, black : max start
  {
    std::vector<std::vector<std::vector<int>>> dp(
        n + 1,
        std::vector<std::vector<int>>(n + 1, std::vector<int>(n + 1, -1)));
    dp[0][0][0] = 0;
    for (int i = 0; i < n; i++) {
      for (int j = 0; j <= n; j++) {
        for (int k = 0; k <= n; k++) {
          if (dp[i][j][k] == -1)
            continue;
          auto go = [&](int i, int j, int k, int val) {
            dp[i][j][k] = std::max(dp[i][j][k], val);
          };
          int stock = dp[i][j][k];
          // normal use
          if (i + 1 < n)
            go(i + 2, j + (s[i] == '0'), k + (s[i] == '1'), stock);
          if (i + 1 < n)
            go(i + 2, j + (s[i + 1] == '0'), k + (s[i + 1] == '1'), stock);
          // use stock
          if (stock)
            go(i + 1, j + (s[i] == '0'), k + (s[i] == '1'), stock - 1);
          // stock
          if (i + 1 < n)
            go(i + 2, j, k, stock + 1);
          go(i + 1, j, k, stock);
        }
      }
    }
    dp[n][0][0] = -1; // special
    int sum[n + 1][2];
    memset(sum, 0, sizeof(sum));
    for (int i = 0; i < n; i++)
      sum[i + 1][s[i] - '0']++;
    for (int i = 0; i < n; i++)
      for (int j = 0; j < 2; j++)
        sum[i + 1][j] += sum[i][j];

    for (int i = 0; i <= n; i++)
      for (int j = 0; j <= n; j++)
        for (int k = 0; k <= n; k++)
          if (dp[i][j][k] != -1) {
            int white = j + sum[n][0] - sum[i][0];
            int black = k + sum[n][1] - sum[i][1];
            ok[white][black] = std::max(ok[white][black], i);
          }
  }

  std::reverse(s.begin(), s.end());
  mint dp[n + 1][n + 1][n + 1]; // pos, white, black
  dp[0][0][0] = 1;
  for (int i = 0; i <= n; i++)
    for (int j = 0; j <= n; j++)
      for (int k = 0; k <= n; k++)
        if (dp[i][j][k] != 0) {
          if (j < n)
            dp[i + (i < n && s[i] == '0')][j + 1][k] += dp[i][j][k];
          if (k < n)
            dp[i + (i < n && s[i] == '1')][j][k + 1] += dp[i][j][k];
        }

  mint res = 0;
  for (int i = 0; i <= n; i++)
    for (int j = 0; j <= n; j++) {
      if (ok[i][j] == -1)
        continue;
      int len = n - ok[i][j];
      for (int k = len; k <= n; k++)
        res += dp[k][i][j];
    }
  std::cout << res << std::endl;
  return 0;
}


Error: 0

Fixed submission:

#include <bits/stdc++.h>

int ri() {
  int n;
  scanf("%d", &n);
  return n;
}

template <int mod> struct ModInt {
  int x;
  ModInt() : x(0) {}
  ModInt(int64_t x) : x(x >= 0 ? x % mod : (mod - -x % mod) % mod) {}
  ModInt &operator+=(const ModInt &p) {
    if ((x += p.x) >= mod)
      x -= mod;
    return *this;
  }
  ModInt &operator-=(const ModInt &p) {
    if ((x += mod - p.x) >= mod)
      x -= mod;
    return *this;
  }
  ModInt &operator*=(const ModInt &p) {
    x = (int64_t)x * p.x % mod;
    return *this;
  }
  ModInt &operator/=(const ModInt &p) {
    *this *= p.inverse();
    return *this;
  }
  ModInt &operator^=(int64_t p) {
    ModInt res = 1;
    for (; p; p >>= 1) {
      if (p & 1)
        res *= *this;
      *this *= *this;
    }
    return *this = res;
  }
  ModInt operator-() const { return ModInt(-x); }
  ModInt operator+(const ModInt &p) const { return ModInt(*this) += p; }
  ModInt operator-(const ModInt &p) const { return ModInt(*this) -= p; }
  ModInt operator*(const ModInt &p) const { return ModInt(*this) *= p; }
  ModInt operator/(const ModInt &p) const { return ModInt(*this) /= p; }
  ModInt operator^(int64_t p) const { return ModInt(*this) ^= p; }
  bool operator==(const ModInt &p) const { return x == p.x; }
  bool operator!=(const ModInt &p) const { return x != p.x; }
  explicit operator int() const { return x; }
  ModInt &operator=(const int p) {
    x = p;
    return *this;
  }
  ModInt inverse() const {
    int a = x, b = mod, u = 1, v = 0, t;
    while (b > 0) {
      t = a / b;
      a -= t * b;
      std::swap(a, b);
      u -= t * v;
      std::swap(u, v);
    }
    return ModInt(u);
  }
  friend std::ostream &operator<<(std::ostream &stream, const ModInt<mod> &p) {
    return stream << p.x;
  }
  friend std::istream &operator>>(std::istream &stream, ModInt<mod> &a) {
    int64_t x;
    stream >> x;
    a = ModInt<mod>(x);
    return stream;
  }
};

template <int mod> struct MComb {
  using mint = ModInt<mod>;
  std::vector<mint> fact;
  std::vector<mint> inv;
  MComb(int n) { // O(n + log(mod))
    fact = std::vector<mint>(n + 1, 1);
    for (int i = 1; i <= n; i++)
      fact[i] = fact[i - 1] * mint(i);
    inv.resize(n + 1);
    inv[n] = fact[n] ^ (mod - 2);
    for (int i = n; i--;)
      inv[i] = inv[i + 1] * mint(i + 1);
  }
  mint ncr(int n, int r) { return fact[n] * inv[r] * inv[n - r]; }
  mint npr(int n, int r) { return fact[n] * inv[n - r]; }
  mint nhr(int n, int r) {
    assert(n + r - 1 < (int)fact.size());
    return ncr(n + r - 1, r);
  }
};

typedef ModInt<998244353> mint;

int main() {
  std::string s;
  std::cin >> s;
  int n = s.size();

  assert(std::count(s.begin(), s.end(), '0') >= n / 10 &&
         std::count(s.begin(), s.end(), '1') >= n / 10);

  std::vector<std::vector<int>> ok(
      n + 1, std::vector<int>(n + 1, -1)); // white, black : max start
  {
    std::vector<std::vector<std::vector<int>>> dp(
        n + 1,
        std::vector<std::vector<int>>(n + 1, std::vector<int>(n + 1, -1)));
    dp[0][0][0] = 0;
    for (int i = 0; i < n; i++) {
      for (int j = 0; j <= n; j++) {
        for (int k = 0; k <= n; k++) {
          if (dp[i][j][k] == -1)
            continue;
          auto go = [&](int i, int j, int k, int val) {
            dp[i][j][k] = std::max(dp[i][j][k], val);
          };
          int stock = dp[i][j][k];
          // normal use
          if (i + 1 < n)
            go(i + 2, j + (s[i] == '0'), k + (s[i] == '1'), stock);
          if (i + 1 < n)
            go(i + 2, j + (s[i + 1] == '0'), k + (s[i + 1] == '1'), stock);
          // use stock
          if (stock)
            go(i + 1, j + (s[i] == '0'), k + (s[i] == '1'), stock - 1);
          // stock
          if (i + 1 < n)
            go(i + 2, j, k, stock + 1);
        }
      }
    }
    dp[n][0][0] = -1; // special
    int sum[n + 1][2];
    memset(sum, 0, sizeof(sum));
    for (int i = 0; i < n; i++)
      sum[i + 1][s[i] - '0']++;
    for (int i = 0; i < n; i++)
      for (int j = 0; j < 2; j++)
        sum[i + 1][j] += sum[i][j];

    for (int i = 0; i <= n; i++)
      for (int j = 0; j <= n; j++)
        for (int k = 0; k <= n; k++)
          if (dp[i][j][k] != -1) {
            int white = j + sum[n][0] - sum[i][0];
            int black = k + sum[n][1] - sum[i][1];
            ok[white][black] = std::max(ok[white][black], i);
          }
  }

  std::reverse(s.begin(), s.end());
  mint dp[n + 1][n + 1][n + 1]; // pos, white, black
  dp[0][0][0] = 1;
  for (int i = 0; i <= n; i++)
    for (int j = 0; j <= n; j++)
      for (int k = 0; k <= n; k++)
        if (dp[i][j][k] != 0) {
          if (j < n)
            dp[i + (i < n && s[i] == '0')][j + 1][k] += dp[i][j][k];
          if (k < n)
            dp[i + (i < n && s[i] == '1')][j][k + 1] += dp[i][j][k];
        }

  mint res = 0;
  for (int i = 0; i <= n; i++)
    for (int j = 0; j <= n; j++) {
      if (ok[i][j] == -1)
        continue;
      int len = n - ok[i][j];
      // dp[len...][i][j]
      for (int k = len; k <= n; k++)
        res += dp[k][i][j];
    }
  std::cout << res << std::endl;
  return 0;
}