Problem p02635 (C++)

Buggy submission (status: Runtime Error):
/**
 * code generated by JHelper
 * More info: https://github.com/AlexeyDmitriev/JHelper
 * @author majk
 */

#ifndef MAJK_LIB
#define MAJK_LIB

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cmath>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;

#define x first
#define y second
typedef std::pair<int, int> pii;
typedef long long ll;
typedef unsigned long long ull;
typedef unsigned int ui;
typedef pair<ui, ui> puu;

template <typename T, typename U>
std::istream &operator>>(std::istream &i, pair<T, U> &p) {
  i >> p.x >> p.y;
  return i;
}
template <typename T> std::istream &operator>>(std::istream &i, vector<T> &t) {
  for (auto &v : t) {
    i >> v;
  }
  return i;
}
template <typename T, typename U>
std::ostream &operator<<(std::ostream &o, const pair<T, U> &p) {
  o << p.x << ' ' << p.y;
  return o;
}
template <typename T>
std::ostream &operator<<(std::ostream &o, const vector<T> &t) {
  if (t.empty())
    o << '\n';
  for (size_t i = 0; i < t.size(); ++i) {
    o << t[i] << " \n"[i == t.size() - 1];
  }
  return o;
}
template <typename T> using minheap = priority_queue<T, vector<T>, greater<T>>;
template <typename T> using maxheap = priority_queue<T, vector<T>, less<T>>;
ui logceil(ll x) { return x ? 8 * sizeof(ll) - __builtin_clzll(x) : 0; }

namespace std {
template <typename T, typename U> struct hash<pair<T, U>> {
  hash<T> t;
  hash<U> u;
  size_t operator()(const pair<T, U> &p) const {
    return t(p.x) ^ (u(p.y) << 7);
  }
};
} // namespace std
template <typename T, typename F> T bsh(T l, T h, const F &f) {
  T r = -1, m;
  while (l <= h) {
    m = (l + h) / 2;
    if (f(m)) {
      l = m + 1;
      r = m;
    } else {
      h = m - 1;
    }
  }
  return r;
}
template <typename F>
double bshd(double l, double h, const F &f, double p = 1e-9) {
  ui r = 3 + (ui)log2((h - l) / p);
  while (r--) {
    double m = (l + h) / 2;
    if (f(m)) {
      l = m;
    } else {
      h = m;
    }
  }
  return (l + h) / 2;
}
template <typename T, typename F> T bsl(T l, T h, const F &f) {
  T r = -1, m;
  while (l <= h) {
    m = (l + h) / 2;
    if (f(m)) {
      h = m - 1;
      r = m;
    } else {
      l = m + 1;
    }
  }
  return r;
}
template <typename F>
double bsld(double l, double h, const F &f, double p = 1e-9) {
  ui r = 3 + (ui)log2((h - l) / p);
  while (r--) {
    double m = (l + h) / 2;
    if (f(m)) {
      h = m;
    } else {
      l = m;
    }
  }
  return (l + h) / 2;
}
template <typename T> T gcd(T a, T b) {
  if (a < b)
    swap(a, b);
  return b ? gcd(b, a % b) : a;
}

template <typename T> class vector2 : public vector<vector<T>> {
public:
  vector2() {}
  vector2(size_t a, size_t b, T t = T())
      : vector<vector<T>>(a, vector<T>(b, t)) {}
};
template <typename T> class vector3 : public vector<vector2<T>> {
public:
  vector3() {}
  vector3(size_t a, size_t b, size_t c, T t = T())
      : vector<vector2<T>>(a, vector2<T>(b, c, t)) {}
};
template <typename T> class vector4 : public vector<vector3<T>> {
public:
  vector4() {}
  vector4(size_t a, size_t b, size_t c, size_t d, T t = T())
      : vector<vector3<T>>(a, vector3<T>(b, c, d, t)) {}
};
template <typename T> class vector5 : public vector<vector4<T>> {
public:
  vector5() {}
  vector5(size_t a, size_t b, size_t c, size_t d, size_t e, T t = T())
      : vector<vector4<T>>(a, vector4<T>(b, c, d, e, t)) {}
};

#endif
#ifndef MOD_H
#define MOD_H

template <unsigned int N> class Field {
  typedef unsigned int ui;
  typedef unsigned long long ull;
  inline ui pow(ui a, ui p) {
    ui r = 1, e = a;
    while (p) {
      if (p & 1) {
        r = ((ull)r * e) % N;
      }
      e = ((ull)e * e) % N;
      p >>= 1;
    }
    return r;
  }
  /*extended GCD(slow):ll t=0,nt=1,r=N,nr=a;while(nr){ll
   * q=r/nr;t-=q*nt;swap(t,nt);r-=q*nr;swap(r,nr);}assert(r<=1);return(t<0)?t+N:t;*/
  inline ui inv(ui a) { return pow(a, N - 2); }

public:
  inline Field(int x = 0) : v(x < 0 ? N + x : x) {}
  inline Field<N> pow(int p) { return (*this) ^ p; }
  inline Field<N> operator^(int p) { return {(int)pow(v, (ui)p)}; }
  inline Field<N> &operator+=(const Field<N> &o) {
    if (v + o.v >= N)
      v += o.v - N;
    else
      v += o.v;
    return *this;
  }
  inline Field<N> &operator-=(const Field<N> &o) {
    if (v < o.v)
      v -= o.v - N;
    else
      v -= o.v;
    return *this;
  }
  inline Field<N> &operator*=(const Field<N> &o) {
    v = (ull)v * o.v % N;
    return *this;
  }
  inline Field<N> &operator/=(const Field<N> &o) { return *this *= inv(o.v); }
  inline Field<N> operator+(const Field<N> &o) const {
    Field<N> r{*this};
    return r += o;
  }
  inline Field<N> operator-(const Field<N> &o) const {
    Field<N> r{*this};
    return r -= o;
  }
  inline Field<N> operator*(const Field<N> &o) const {
    Field<N> r{*this};
    return r *= o;
  }
  inline Field<N> operator/(const Field<N> &o) const {
    Field<N> r{*this};
    return r /= o;
  }
  inline Field<N> operator-() {
    if (v)
      return {(int)(N - v)};
    else
      return {0};
  };
  inline Field<N> &operator++() {
    ++v;
    if (v == N)
      v = 0;
    return *this;
  }
  inline Field<N> operator++(int) {
    Field<N> r{*this};
    ++*this;
    return r;
  }
  inline Field<N> &operator--() {
    --v;
    if (v == -1)
      v = N - 1;
    return *this;
  }
  inline Field<N> operator--(int) {
    Field<N> r{*this};
    --*this;
    return r;
  }
  inline bool operator==(const Field<N> &o) const { return o.v == v; }
  inline bool operator!=(const Field<N> &o) const { return o.v != v; }
  inline explicit operator ui() const { return v; }
  inline static vector<Field<N>> fact(int t) {
    vector<Field<N>> F(t + 1, 1);
    for (int i = 2; i <= t; ++i) {
      F[i] = F[i - 1] * i;
    }
    return F;
  }
  inline static vector<Field<N>> invfact(int t) {
    vector<Field<N>> F(t + 1, 1);
    Field<N> X{1};
    for (int i = 2; i <= t; ++i) {
      X = X * i;
    }
    F[t] = 1 / X;
    for (int i = t - 1; i >= 2; --i) {
      F[i] = F[i + 1] * (i + 1);
    }
    return F;
  }

private:
  ui v;
};
template <unsigned int N> istream &operator>>(std::istream &is, Field<N> &f) {
  unsigned int v;
  is >> v;
  f = v;
  return is;
}
template <unsigned int N>
ostream &operator<<(std::ostream &os, const Field<N> &f) {
  return os << (unsigned int)f;
}
template <unsigned int N> Field<N> operator+(int i, const Field<N> &f) {
  return Field<N>(i) + f;
}
template <unsigned int N> Field<N> operator-(int i, const Field<N> &f) {
  return Field<N>(i) - f;
}
template <unsigned int N> Field<N> operator*(int i, const Field<N> &f) {
  return Field<N>(i) * f;
}
template <unsigned int N> Field<N> operator/(int i, const Field<N> &f) {
  return Field<N>(i) / f;
}

typedef Field<1000000007> FieldMod;

struct Ring {
  template <typename T> static T div(T p, T q, T N) {
    T t = 0, nt = 1, r = N, nr = q;
    while (nr != 0) {
      T q = r / nr;
      t -= q * nt;
      r -= q * nr;
      swap(t, nt);
      swap(r, nr);
    }
    t = (t < 0) ? t + N : t;
    r = (r < 0) ? r + N : r;
    if (gcd(p, N) % r != 0) {
      return 0;
    }
    return (t * p / r) % N;
  }
};
#endif

typedef Field<998244353> FF;

class CShift {
public:
  void solve(istream &cin, ostream &cout) {
    string S;
    cin >> S;
    int K;
    cin >> K;
    int N = S.size();

    vector<int> A;
    int cnt = 0;
    for (int i = 0; i < N; ++i) {
      if (S[i] == '0') {
        A.push_back(cnt);
        cnt = 0;
      } else
        cnt++;
    }
    A.push_back(cnt);
    reverse(A.begin(), A.end());
    int M = A.size();

    vector3<FF> DP(M + 1, K + 1, K + 1, 0);
    DP[0][0][0] = 1;
    for (int i = 0; i < M; ++i) {
      int a = A[i];
      for (int j = 0; j <= K; ++j) {
        vector<FF> Pref(K + 2, 0);
        for (int k = 0; k <= K; ++k)
          Pref[k + 1] = Pref[k] + DP[i][j][k];

        for (int k = 0; k <= j; ++k) {
          for (int l = 1; l <= a; ++l) {
            if (l + j <= K) {
              DP[i + 1][j + l][k] += DP[i][j][k]; // store ones
            }
          }

          DP[i + 1][j][k] += Pref[k + 1];
        }
      }
    }

    FF ans = 0;
    for (int i = 0; i <= K; ++i) {
      ans += DP[M][i][i];
    }
    cout << ans << '\n';
  }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  cout.tie(nullptr);
  CShift solver;
  std::istream &in(std::cin);
  std::ostream &out(std::cout);
  solver.solve(in, out);
  return 0;
}


Fixed submission:
/**
 * code generated by JHelper
 * More info: https://github.com/AlexeyDmitriev/JHelper
 * @author majk
 */

#ifndef MAJK_LIB
#define MAJK_LIB

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cmath>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;

#define x first
#define y second
typedef std::pair<int, int> pii;
typedef long long ll;
typedef unsigned long long ull;
typedef unsigned int ui;
typedef pair<ui, ui> puu;

template <typename T, typename U>
std::istream &operator>>(std::istream &i, pair<T, U> &p) {
  i >> p.x >> p.y;
  return i;
}
template <typename T> std::istream &operator>>(std::istream &i, vector<T> &t) {
  for (auto &v : t) {
    i >> v;
  }
  return i;
}
template <typename T, typename U>
std::ostream &operator<<(std::ostream &o, const pair<T, U> &p) {
  o << p.x << ' ' << p.y;
  return o;
}
template <typename T>
std::ostream &operator<<(std::ostream &o, const vector<T> &t) {
  if (t.empty())
    o << '\n';
  for (size_t i = 0; i < t.size(); ++i) {
    o << t[i] << " \n"[i == t.size() - 1];
  }
  return o;
}
template <typename T> using minheap = priority_queue<T, vector<T>, greater<T>>;
template <typename T> using maxheap = priority_queue<T, vector<T>, less<T>>;
ui logceil(ll x) { return x ? 8 * sizeof(ll) - __builtin_clzll(x) : 0; }

namespace std {
template <typename T, typename U> struct hash<pair<T, U>> {
  hash<T> t;
  hash<U> u;
  size_t operator()(const pair<T, U> &p) const {
    return t(p.x) ^ (u(p.y) << 7);
  }
};
} // namespace std
template <typename T, typename F> T bsh(T l, T h, const F &f) {
  T r = -1, m;
  while (l <= h) {
    m = (l + h) / 2;
    if (f(m)) {
      l = m + 1;
      r = m;
    } else {
      h = m - 1;
    }
  }
  return r;
}
template <typename F>
double bshd(double l, double h, const F &f, double p = 1e-9) {
  ui r = 3 + (ui)log2((h - l) / p);
  while (r--) {
    double m = (l + h) / 2;
    if (f(m)) {
      l = m;
    } else {
      h = m;
    }
  }
  return (l + h) / 2;
}
template <typename T, typename F> T bsl(T l, T h, const F &f) {
  T r = -1, m;
  while (l <= h) {
    m = (l + h) / 2;
    if (f(m)) {
      h = m - 1;
      r = m;
    } else {
      l = m + 1;
    }
  }
  return r;
}
template <typename F>
double bsld(double l, double h, const F &f, double p = 1e-9) {
  ui r = 3 + (ui)log2((h - l) / p);
  while (r--) {
    double m = (l + h) / 2;
    if (f(m)) {
      h = m;
    } else {
      l = m;
    }
  }
  return (l + h) / 2;
}
template <typename T> T gcd(T a, T b) {
  if (a < b)
    swap(a, b);
  return b ? gcd(b, a % b) : a;
}

template <typename T> class vector2 : public vector<vector<T>> {
public:
  vector2() {}
  vector2(size_t a, size_t b, T t = T())
      : vector<vector<T>>(a, vector<T>(b, t)) {}
};
template <typename T> class vector3 : public vector<vector2<T>> {
public:
  vector3() {}
  vector3(size_t a, size_t b, size_t c, T t = T())
      : vector<vector2<T>>(a, vector2<T>(b, c, t)) {}
};
template <typename T> class vector4 : public vector<vector3<T>> {
public:
  vector4() {}
  vector4(size_t a, size_t b, size_t c, size_t d, T t = T())
      : vector<vector3<T>>(a, vector3<T>(b, c, d, t)) {}
};
template <typename T> class vector5 : public vector<vector4<T>> {
public:
  vector5() {}
  vector5(size_t a, size_t b, size_t c, size_t d, size_t e, T t = T())
      : vector<vector4<T>>(a, vector4<T>(b, c, d, e, t)) {}
};

#endif
#ifndef MOD_H
#define MOD_H

template <unsigned int N> class Field {
  typedef unsigned int ui;
  typedef unsigned long long ull;
  inline ui pow(ui a, ui p) {
    ui r = 1, e = a;
    while (p) {
      if (p & 1) {
        r = ((ull)r * e) % N;
      }
      e = ((ull)e * e) % N;
      p >>= 1;
    }
    return r;
  }
  /*extended GCD(slow):ll t=0,nt=1,r=N,nr=a;while(nr){ll
   * q=r/nr;t-=q*nt;swap(t,nt);r-=q*nr;swap(r,nr);}assert(r<=1);return(t<0)?t+N:t;*/
  inline ui inv(ui a) { return pow(a, N - 2); }

public:
  inline Field(int x = 0) : v(x < 0 ? N + x : x) {}
  inline Field<N> pow(int p) { return (*this) ^ p; }
  inline Field<N> operator^(int p) { return {(int)pow(v, (ui)p)}; }
  inline Field<N> &operator+=(const Field<N> &o) {
    if (v + o.v >= N)
      v += o.v - N;
    else
      v += o.v;
    return *this;
  }
  inline Field<N> &operator-=(const Field<N> &o) {
    if (v < o.v)
      v -= o.v - N;
    else
      v -= o.v;
    return *this;
  }
  inline Field<N> &operator*=(const Field<N> &o) {
    v = (ull)v * o.v % N;
    return *this;
  }
  inline Field<N> &operator/=(const Field<N> &o) { return *this *= inv(o.v); }
  inline Field<N> operator+(const Field<N> &o) const {
    Field<N> r{*this};
    return r += o;
  }
  inline Field<N> operator-(const Field<N> &o) const {
    Field<N> r{*this};
    return r -= o;
  }
  inline Field<N> operator*(const Field<N> &o) const {
    Field<N> r{*this};
    return r *= o;
  }
  inline Field<N> operator/(const Field<N> &o) const {
    Field<N> r{*this};
    return r /= o;
  }
  inline Field<N> operator-() {
    if (v)
      return {(int)(N - v)};
    else
      return {0};
  };
  inline Field<N> &operator++() {
    ++v;
    if (v == N)
      v = 0;
    return *this;
  }
  inline Field<N> operator++(int) {
    Field<N> r{*this};
    ++*this;
    return r;
  }
  inline Field<N> &operator--() {
    --v;
    if (v == -1)
      v = N - 1;
    return *this;
  }
  inline Field<N> operator--(int) {
    Field<N> r{*this};
    --*this;
    return r;
  }
  inline bool operator==(const Field<N> &o) const { return o.v == v; }
  inline bool operator!=(const Field<N> &o) const { return o.v != v; }
  inline explicit operator ui() const { return v; }
  inline static vector<Field<N>> fact(int t) {
    vector<Field<N>> F(t + 1, 1);
    for (int i = 2; i <= t; ++i) {
      F[i] = F[i - 1] * i;
    }
    return F;
  }
  inline static vector<Field<N>> invfact(int t) {
    vector<Field<N>> F(t + 1, 1);
    Field<N> X{1};
    for (int i = 2; i <= t; ++i) {
      X = X * i;
    }
    F[t] = 1 / X;
    for (int i = t - 1; i >= 2; --i) {
      F[i] = F[i + 1] * (i + 1);
    }
    return F;
  }

private:
  ui v;
};
template <unsigned int N> istream &operator>>(std::istream &is, Field<N> &f) {
  unsigned int v;
  is >> v;
  f = v;
  return is;
}
template <unsigned int N>
ostream &operator<<(std::ostream &os, const Field<N> &f) {
  return os << (unsigned int)f;
}
template <unsigned int N> Field<N> operator+(int i, const Field<N> &f) {
  return Field<N>(i) + f;
}
template <unsigned int N> Field<N> operator-(int i, const Field<N> &f) {
  return Field<N>(i) - f;
}
template <unsigned int N> Field<N> operator*(int i, const Field<N> &f) {
  return Field<N>(i) * f;
}
template <unsigned int N> Field<N> operator/(int i, const Field<N> &f) {
  return Field<N>(i) / f;
}

typedef Field<1000000007> FieldMod;

struct Ring {
  template <typename T> static T div(T p, T q, T N) {
    T t = 0, nt = 1, r = N, nr = q;
    while (nr != 0) {
      T q = r / nr;
      t -= q * nt;
      r -= q * nr;
      swap(t, nt);
      swap(r, nr);
    }
    t = (t < 0) ? t + N : t;
    r = (r < 0) ? r + N : r;
    if (gcd(p, N) % r != 0) {
      return 0;
    }
    return (t * p / r) % N;
  }
};
#endif

typedef Field<998244353> FF;

class CShift {
public:
  void solve(istream &cin, ostream &cout) {
    string S;
    cin >> S;
    int K;
    cin >> K;
    int N = S.size();
    K = min(K, N);

    vector<int> A;
    int cnt = 0;
    for (int i = 0; i < N; ++i) {
      if (S[i] == '0') {
        A.push_back(cnt);
        cnt = 0;
      } else
        cnt++;
    }
    A.push_back(cnt);
    reverse(A.begin(), A.end());
    int M = A.size();

    vector3<FF> DP(M + 1, K + 1, K + 1, 0);
    DP[0][0][0] = 1;
    for (int i = 0; i < M; ++i) {
      int a = A[i];
      for (int j = 0; j <= K; ++j) {
        vector<FF> Pref(K + 2, 0);
        for (int k = 0; k <= K; ++k)
          Pref[k + 1] = Pref[k] + DP[i][j][k];

        for (int k = 0; k <= j; ++k) {
          for (int l = 1; l <= a; ++l) {
            if (l + j <= K) {
              DP[i + 1][j + l][k] += DP[i][j][k]; // store ones
            }
          }

          DP[i + 1][j][k] += Pref[k + 1];
        }
      }
    }

    FF ans = 0;
    for (int i = 0; i <= K; ++i) {
      ans += DP[M][i][i];
    }
    cout << ans << '\n';
  }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  cout.tie(nullptr);
  CShift solver;
  std::istream &in(std::cin);
  std::ostream &out(std::cout);
  solver.solve(in, out);
  return 0;
}