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

#include <bits/stdc++.h>

using namespace std;

#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (int)(n)-1; i >= 0; --i)
#define FOR(i, a, n) for (int i = (a); i < (n); ++i)
#define RFOR(i, a, n) for (int i = (int)(n)-1; i >= (a); --i)

#define SZ(x) ((int)(x).size())
#define ALL(x) (x).begin(), (x).end()

#define DUMP(x) cerr << #x << " = " << (x) << endl
#define DEBUG(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  REP(i, SZ(v)) {
    if (i)
      os << ", ";
    os << v[i];
  }
  return os << "]";
}

template <class T, class U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  return os << "(" << p.first << " " << p.second << ")";
}

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

using ll = long long;
using ull = unsigned long long;
using ld = long double;
using P = pair<int, int>;
// using vi = vector<int>;
// using vll = vector<ll>;
// using vvi = vector<vi>;
// using vvll = vector<vll>;

const ll MOD = 1e9 + 7;
const int INF = INT_MAX / 2;
const ll LINF = LLONG_MAX / 2;
const ld eps = 1e-9;

int gcd(int a, int b) {
  if (a == 0)
    return b;
  return gcd(b % a, a);
}

template <int64_t mod> class modint {
  int64_t x;

public:
  modint(int64_t x = 0) : x(x < 0 ? ((x % mod) + mod) % mod : x % mod) {}

  const modint operator-() const { return x == 0 ? 0 : mod - x; }

  modint &operator+=(const modint &rhs) {
    if ((x += rhs.x) >= mod)
      x -= mod;
    return *this;
  }
  modint &operator-=(const modint &rhs) { return *this += -rhs; }
  modint &operator*=(const modint &rhs) {
    (x *= rhs.x) %= mod;
    return *this;
  }
  modint &operator/=(const modint &rhs) { return *this *= rhs.pow(mod - 2); }

  friend const modint operator+(modint lhs, const modint &rhs) {
    return lhs += rhs;
  }
  friend const modint operator-(modint lhs, const modint &rhs) {
    return lhs -= rhs;
  }
  friend const modint operator*(modint lhs, const modint &rhs) {
    return lhs *= rhs;
  }
  friend const modint operator/(modint lhs, const modint &rhs) {
    return lhs /= rhs;
  }

  const modint pow(int64_t n) const {
    modint ret = 1, tmp = *this;
    while (n) {
      if (n & 1)
        ret *= tmp;
      tmp *= tmp;
      n >>= 1;
    }
    return ret;
  }

  friend bool operator==(const modint &lhs, const modint &rhs) {
    return lhs.x == rhs.x;
  }
  friend bool operator!=(const modint &lhs, const modint &rhs) {
    return lhs.x != rhs.x;
  }

  friend ostream &operator<<(ostream &os, const modint &a) { return os << a.x; }
  friend istream &operator>>(istream &is, modint &a) {
    int64_t tmp;
    is >> tmp;
    a = tmp;
    return is;
  }
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);

  // ifstream in("in.txt");
  // cin.rdbuf(in.rdbuf());

  string S;
  cin >> S;
  int K;
  cin >> K;

  using Int = modint<998244353>;

  vector<int> A;
  {
    int cnt = 0;
    for (char c : S) {
      if (c == '0')
        ++cnt;
      else if (cnt > 0)
        A.push_back(cnt);
    }
  }

  if (A.empty()) {
    cout << 1 << endl;
    return 0;
  }

  sort(ALL(A));

  int s = A.size();
  unordered_map<int, int> cntA;
  for (int a : A)
    ++cntA[a];

  int maxA = *max_element(ALL(A));
  int maxCnt = 0;
  for (auto &tp : cntA)
    chmax(maxCnt, tp.second);

  using vi = vector<Int>;
  using vvi = vector<vi>;
  using vvvi = vector<vvi>;
  using vvvvi = vector<vvvi>;

  // vvvvi dp(s+1, vvvi(s+1, vvi(maxA+1, vi(maxCnt+1))));
  vvvvi dp(s + 1, vvvi(1, vvi(1, vi(1))));

  dp[0][0][0][0] = 1;
  for (int i = 0; i < s; ++i) {
    dp[i + 1].resize(i + 2);

    for (int j = 0; j <= i + 1; ++j) {
      dp[i + 1][j].resize(A[i] + 1);

      for (int a = 0; a <= A[i]; ++a) {
        dp[i + 1][j][a].resize(cntA[a] + 1);
      }
    }
  }

  for (int i = 0; i < s; ++i) {

    for (int j = 0; j <= i; ++j) {

      for (int a = 0; a <= A[i]; ++a) {

        for (int b = 0; b <= min(i > 0 ? A[i - 1] : 0, a); ++b) {

          if (b == a) {
            for (int p = 0; p <= cntA[a]; ++p) {

              if (p + 1 <= cntA[a]) {
                dp[i + 1][j][a][p + 1] += dp[i][j][b][p];
              } else if (j + 1 <= i + 1) {
                dp[i + 1][j + 1][a][p] += dp[i][j][b][p];
              }
            }
          } else {
            if (cntA[a] > 0) {
              dp[i + 1][j][a][1] += accumulate(ALL(dp[i][j][b]), Int{0});
            } else if (j + 1 <= i + 1) {
              dp[i + 1][j + 1][a][0] += accumulate(ALL(dp[i][j][b]), Int{0});
            }
          }
        }
      }
    }
  }

  Int ans = 0;

  for (int k = 0; k <= min(s, K); ++k) {
    for (int a = 0; a <= maxA; ++a) {
      ans += accumulate(ALL(dp[s][k][a]), Int{0});
    }
  }

  cout << ans << endl;

  return 0;
}


Error: 0

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (int)(n)-1; i >= 0; --i)
#define FOR(i, a, n) for (int i = (a); i < (n); ++i)
#define RFOR(i, a, n) for (int i = (int)(n)-1; i >= (a); --i)

#define SZ(x) ((int)(x).size())
#define ALL(x) (x).begin(), (x).end()

#define DUMP(x) cerr << #x << " = " << (x) << endl
#define DEBUG(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  REP(i, SZ(v)) {
    if (i)
      os << ", ";
    os << v[i];
  }
  return os << "]";
}

template <class T, class U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  return os << "(" << p.first << " " << p.second << ")";
}

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

using ll = long long;
using ull = unsigned long long;
using ld = long double;
using P = pair<int, int>;
// using vi = vector<int>;
// using vll = vector<ll>;
// using vvi = vector<vi>;
// using vvll = vector<vll>;

const ll MOD = 1e9 + 7;
const int INF = INT_MAX / 2;
const ll LINF = LLONG_MAX / 2;
const ld eps = 1e-9;

int gcd(int a, int b) {
  if (a == 0)
    return b;
  return gcd(b % a, a);
}

template <int64_t mod> class modint {
  int64_t x;

public:
  modint(int64_t x = 0) : x(x < 0 ? ((x % mod) + mod) % mod : x % mod) {}

  const modint operator-() const { return x == 0 ? 0 : mod - x; }

  modint &operator+=(const modint &rhs) {
    if ((x += rhs.x) >= mod)
      x -= mod;
    return *this;
  }
  modint &operator-=(const modint &rhs) { return *this += -rhs; }
  modint &operator*=(const modint &rhs) {
    (x *= rhs.x) %= mod;
    return *this;
  }
  modint &operator/=(const modint &rhs) { return *this *= rhs.pow(mod - 2); }

  friend const modint operator+(modint lhs, const modint &rhs) {
    return lhs += rhs;
  }
  friend const modint operator-(modint lhs, const modint &rhs) {
    return lhs -= rhs;
  }
  friend const modint operator*(modint lhs, const modint &rhs) {
    return lhs *= rhs;
  }
  friend const modint operator/(modint lhs, const modint &rhs) {
    return lhs /= rhs;
  }

  const modint pow(int64_t n) const {
    modint ret = 1, tmp = *this;
    while (n) {
      if (n & 1)
        ret *= tmp;
      tmp *= tmp;
      n >>= 1;
    }
    return ret;
  }

  friend bool operator==(const modint &lhs, const modint &rhs) {
    return lhs.x == rhs.x;
  }
  friend bool operator!=(const modint &lhs, const modint &rhs) {
    return lhs.x != rhs.x;
  }

  friend ostream &operator<<(ostream &os, const modint &a) { return os << a.x; }
  friend istream &operator>>(istream &is, modint &a) {
    int64_t tmp;
    is >> tmp;
    a = tmp;
    return is;
  }
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);

  // ifstream in("in.txt");
  // cin.rdbuf(in.rdbuf());

  string S;
  cin >> S;
  int K;
  cin >> K;

  using Int = modint<998244353>;

  vector<int> A;
  {
    int cnt = 0;
    for (char c : S) {
      if (c == '0')
        ++cnt;
      else if (cnt > 0)
        A.push_back(cnt);
    }
  }

  if (A.empty()) {
    cout << 1 << endl;
    return 0;
  }

  sort(ALL(A));

  int s = A.size();
  unordered_map<int, int> cntA;
  for (int a : A)
    ++cntA[a];

  int maxA = *max_element(ALL(A));
  int maxCnt = 0;
  for (auto &tp : cntA)
    chmax(maxCnt, tp.second);

  using vi = vector<Int>;
  using vvi = vector<vi>;
  using vvvi = vector<vvi>;
  using vvvvi = vector<vvvi>;

  vvvvi dp(s + 1, vvvi(s + 1, vvi(maxA + 1, vi(maxCnt + 1))));

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

  for (int i = 0; i < s; ++i) {

    for (int j = 0; j <= i; ++j) {

      for (int a = 0; a <= A[i]; ++a) {

        for (int b = 0; b <= min(i > 0 ? A[i - 1] : 0, a); ++b) {

          if (b == a) {
            for (int p = 0; p <= cntA[a]; ++p) {

              if (p + 1 <= cntA[a]) {
                dp[i + 1][j][a][p + 1] += dp[i][j][b][p];
              } else if (j + 1 <= i + 1) {
                dp[i + 1][j + 1][a][p] += dp[i][j][b][p];
              }
            }
          } else {
            if (cntA[a] > 0) {
              dp[i + 1][j][a][1] += accumulate(ALL(dp[i][j][b]), Int{0});
            } else if (j + 1 <= i + 1) {
              dp[i + 1][j + 1][a][0] += accumulate(ALL(dp[i][j][b]), Int{0});
            }
          }
        }
      }
    }
  }

  Int ans = 0;

  for (int k = 0; k <= min(s, K); ++k) {
    for (int a = 0; a <= maxA; ++a) {
      ans += accumulate(ALL(dp[s][k][a]), Int{0});
    }
  }

  cout << ans << endl;

  return 0;
}