Problem p02635 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>

using i32 = int_fast32_t;
using i64 = int_fast64_t;
using u32 = uint_fast32_t;
using u64 = uint_fast64_t;
using f64 = double;
using f80 = long double;

#define FOR(v, a, b) for (i64 v = (a); v < (b); ++v)
#define FORE(v, a, b) for (i64 v = (a); v <= (b); ++v)
#define REP(v, n) FOR(v, 0, n)
#define REPE(v, n) FORE(v, 0, n)
#define REV(v, a, b) for (i64 v = (a); v >= (b); --v)
#define ALL(x) (x).begin(), (x).end()
#define RALL(x) (x).rbegin(), (x).rend()
#define ITR(it, c) for (auto it = (c).begin(); it != (c).end(); ++it)
#define RITR(it, c) for (auto it = (c).rbegin(); it != (c).rend(); ++it)
#define EXIST(c, x) ((c).find(x) != (c).end())
#define fst first
#define snd second
#define UNIQ(v) (v).erase(unique(ALL(v)), (v).end())
#define bit(i) (1LL << (i))

#ifdef DEBUG
#include <Mylib/Debug/debug.cpp>
#else
#define dump(...) ((void)0)
#endif

using namespace std;

template <typename I> void join(ostream &ost, I s, I t, string d = " ") {
  for (auto i = s; i != t; ++i) {
    if (i != s)
      ost << d;
    ost << *i;
  }
  ost << endl;
}
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &a : v)
    is >> a;
  return is;
}

template <typename T, typename U> bool chmin(T &a, const U &b) {
  return (a > b ? a = b, true : false);
}
template <typename T, typename U> bool chmax(T &a, const U &b) {
  return (a < b ? a = b, true : false);
}

template <typename T, size_t N, typename U>
void fill_array(T (&a)[N], const U &v) {
  fill((U *)a, (U *)(a + N), v);
}
template <typename T> auto make_vector(int n, int m, const T &value) {
  return vector<vector<T>>(n, vector<T>(m, value));
}

template <typename T> void print(const T &value) { std::cout << value << "\\n"; }
template <typename T, typename... Args>
void print(const T &value, const Args &...args) {
  std::cout << value << " ";
  print(args...);
}

struct Init {
  Init() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(12);
    cerr << fixed << setprecision(12);
  }
} init;

/**
 * @title Modint
 * @docs mint.md
 */
template <uint32_t M> class ModInt {
public:
  constexpr static uint32_t MOD = M;
  uint64_t val;

  constexpr ModInt() : val(0) {}
  constexpr ModInt(int64_t n) {
    if (n >= M)
      val = n % M;
    else if (n < 0)
      val = n % M + M;
    else
      val = n;
  }

  inline constexpr auto operator+(const ModInt &a) const {
    return ModInt(val + a.val);
  }
  inline constexpr auto operator-(const ModInt &a) const {
    return ModInt(val - a.val);
  }
  inline constexpr auto operator*(const ModInt &a) const {
    return ModInt(val * a.val);
  }
  inline constexpr auto operator/(const ModInt &a) const {
    return ModInt(val * a.inv().val);
  }

  inline constexpr auto &operator=(const ModInt &a) {
    val = a.val;
    return *this;
  }
  inline constexpr auto &operator+=(const ModInt &a) {
    if ((val += a.val) >= M)
      val -= M;
    return *this;
  }
  inline constexpr auto &operator-=(const ModInt &a) {
    if (val < a.val)
      val += M;
    val -= a.val;
    return *this;
  }
  inline constexpr auto &operator*=(const ModInt &a) {
    (val *= a.val) %= M;
    return *this;
  }
  inline constexpr auto &operator/=(const ModInt &a) {
    (val *= a.inv().val) %= M;
    return *this;
  }

  inline constexpr bool operator==(const ModInt &a) const {
    return val == a.val;
  }
  inline constexpr bool operator!=(const ModInt &a) const {
    return val != a.val;
  }

  inline constexpr auto &operator++() {
    *this += 1;
    return *this;
  }
  inline constexpr auto &operator--() {
    *this -= 1;
    return *this;
  }

  inline constexpr auto operator++(int) {
    auto t = *this;
    *this += 1;
    return t;
  }
  inline constexpr auto operator--(int) {
    auto t = *this;
    *this -= 1;
    return t;
  }

  inline constexpr static ModInt power(int64_t n, int64_t p) {
    if (p < 0)
      return power(n, -p).inv();

    int64_t ret = 1, e = n % M;
    for (; p; (e *= e) %= M, p >>= 1)
      if (p & 1)
        (ret *= e) %= M;
    return ret;
  }

  inline constexpr static ModInt inv(int64_t a) {
    int64_t b = M, u = 1, v = 0;

    while (b) {
      int64_t t = a / b;
      a -= t * b;
      std::swap(a, b);
      u -= t * v;
      std::swap(u, v);
    }

    u %= M;
    if (u < 0)
      u += M;

    return u;
  }

  inline constexpr static auto frac(int64_t a, int64_t b) {
    return ModInt(a) / ModInt(b);
  }

  inline constexpr auto power(int64_t p) const { return power(val, p); }
  inline constexpr auto inv() const { return inv(val); }

  friend inline constexpr auto operator-(const ModInt &a) {
    return ModInt(-a.val);
  }

  friend inline constexpr auto operator+(int64_t a, const ModInt &b) {
    return ModInt(a) + b;
  }
  friend inline constexpr auto operator-(int64_t a, const ModInt &b) {
    return ModInt(a) - b;
  }
  friend inline constexpr auto operator*(int64_t a, const ModInt &b) {
    return ModInt(a) * b;
  }
  friend inline constexpr auto operator/(int64_t a, const ModInt &b) {
    return ModInt(a) / b;
  }

  friend std::istream &operator>>(std::istream &s, ModInt<M> &a) {
    s >> a.val;
    return s;
  }
  friend std::ostream &operator<<(std::ostream &s, const ModInt<M> &a) {
    s << a.val;
    return s;
  }

  template <int N> inline static auto div() {
    static auto value = inv(N);
    return value;
  }

  explicit operator int32_t() const noexcept { return val; }
  explicit operator int64_t() const noexcept { return val; }
};

using mint = ModInt<998244353>;

/**
 * @title Run length encoding
 * @docs run_length_encoding.md
 */
template <typename Container, typename T = typename Container::value_type>
auto run_length_encoding(const Container &v) {
  std::vector<std::pair<T, int64_t>> ret;

  for (auto &x : v) {
    if (ret.empty())
      ret.push_back({x, 1});
    else if (ret.back().first == x)
      ++ret.back().second;
    else
      ret.push_back({x, 1});
  }

  return ret;
}

mint dp[2][301][301];

int main() {
  string S;
  int K;

  while (cin >> S >> K) {
    const int L = count(ALL(S), '1');

    auto r = run_length_encoding(S);
    // dump(r);

    vector<int> a;

    // dump(r);
    if (S[0] == '0')
      a.push_back(0);

    REP(i, r.size()) {
      if (r[i].fst == '1')
        a.push_back(r[i].snd);
      else {
        REP(j, r[i].snd - 1) { a.push_back(0); }
      }
    }

    const int N = a.size();

    dump(a);

    fill_array(dp, 0);

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

    REP(i, N) {
      int cur = i % 2;
      int nxt = (i + 1) % 2;

      FORE(j, 0, L) {
        FORE(k, 0, L) {
          FORE(d, max<int>(-j, -a[i]), min<int>(L - j, L - a[i])) {
            if (k + max<int>(0, d) <= L) {
              dp[nxt][j + d][k + max<int>(0, d)] += dp[cur][j][k];
            }
          }
        }
      }

      FORE(j, 0, L) {
        FORE(k, 0, L) { dp[cur][j][k] = 0; }
      }
    }

    mint ans = 0;

    FORE(i, 0, min(K, L)) { ans += dp[N % 2][0][i]; }

    print(ans);
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using i32 = int_fast32_t;
using i64 = int_fast64_t;
using u32 = uint_fast32_t;
using u64 = uint_fast64_t;
using f64 = double;
using f80 = long double;

#define FOR(v, a, b) for (i64 v = (a); v < (b); ++v)
#define FORE(v, a, b) for (i64 v = (a); v <= (b); ++v)
#define REP(v, n) FOR(v, 0, n)
#define REPE(v, n) FORE(v, 0, n)
#define REV(v, a, b) for (i64 v = (a); v >= (b); --v)
#define ALL(x) (x).begin(), (x).end()
#define RALL(x) (x).rbegin(), (x).rend()
#define ITR(it, c) for (auto it = (c).begin(); it != (c).end(); ++it)
#define RITR(it, c) for (auto it = (c).rbegin(); it != (c).rend(); ++it)
#define EXIST(c, x) ((c).find(x) != (c).end())
#define fst first
#define snd second
#define UNIQ(v) (v).erase(unique(ALL(v)), (v).end())
#define bit(i) (1LL << (i))

#ifdef DEBUG
#include <Mylib/Debug/debug.cpp>
#else
#define dump(...) ((void)0)
#endif

using namespace std;

template <typename I> void join(ostream &ost, I s, I t, string d = " ") {
  for (auto i = s; i != t; ++i) {
    if (i != s)
      ost << d;
    ost << *i;
  }
  ost << endl;
}
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &a : v)
    is >> a;
  return is;
}

template <typename T, typename U> bool chmin(T &a, const U &b) {
  return (a > b ? a = b, true : false);
}
template <typename T, typename U> bool chmax(T &a, const U &b) {
  return (a < b ? a = b, true : false);
}

template <typename T, size_t N, typename U>
void fill_array(T (&a)[N], const U &v) {
  fill((U *)a, (U *)(a + N), v);
}
template <typename T> auto make_vector(int n, int m, const T &value) {
  return vector<vector<T>>(n, vector<T>(m, value));
}

template <typename T> void print(const T &value) { std::cout << value << "\\n"; }
template <typename T, typename... Args>
void print(const T &value, const Args &...args) {
  std::cout << value << " ";
  print(args...);
}

struct Init {
  Init() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(12);
    cerr << fixed << setprecision(12);
  }
} init;

/**
 * @title Modint
 * @docs mint.md
 */
template <uint32_t M> class ModInt {
public:
  constexpr static uint32_t MOD = M;
  uint64_t val;

  constexpr ModInt() : val(0) {}
  constexpr ModInt(int64_t n) {
    if (n >= M)
      val = n % M;
    else if (n < 0)
      val = n % M + M;
    else
      val = n;
  }

  inline constexpr auto operator+(const ModInt &a) const {
    return ModInt(val + a.val);
  }
  inline constexpr auto operator-(const ModInt &a) const {
    return ModInt(val - a.val);
  }
  inline constexpr auto operator*(const ModInt &a) const {
    return ModInt(val * a.val);
  }
  inline constexpr auto operator/(const ModInt &a) const {
    return ModInt(val * a.inv().val);
  }

  inline constexpr auto &operator=(const ModInt &a) {
    val = a.val;
    return *this;
  }
  inline constexpr auto &operator+=(const ModInt &a) {
    if ((val += a.val) >= M)
      val -= M;
    return *this;
  }
  inline constexpr auto &operator-=(const ModInt &a) {
    if (val < a.val)
      val += M;
    val -= a.val;
    return *this;
  }
  inline constexpr auto &operator*=(const ModInt &a) {
    (val *= a.val) %= M;
    return *this;
  }
  inline constexpr auto &operator/=(const ModInt &a) {
    (val *= a.inv().val) %= M;
    return *this;
  }

  inline constexpr bool operator==(const ModInt &a) const {
    return val == a.val;
  }
  inline constexpr bool operator!=(const ModInt &a) const {
    return val != a.val;
  }

  inline constexpr auto &operator++() {
    *this += 1;
    return *this;
  }
  inline constexpr auto &operator--() {
    *this -= 1;
    return *this;
  }

  inline constexpr auto operator++(int) {
    auto t = *this;
    *this += 1;
    return t;
  }
  inline constexpr auto operator--(int) {
    auto t = *this;
    *this -= 1;
    return t;
  }

  inline constexpr static ModInt power(int64_t n, int64_t p) {
    if (p < 0)
      return power(n, -p).inv();

    int64_t ret = 1, e = n % M;
    for (; p; (e *= e) %= M, p >>= 1)
      if (p & 1)
        (ret *= e) %= M;
    return ret;
  }

  inline constexpr static ModInt inv(int64_t a) {
    int64_t b = M, u = 1, v = 0;

    while (b) {
      int64_t t = a / b;
      a -= t * b;
      std::swap(a, b);
      u -= t * v;
      std::swap(u, v);
    }

    u %= M;
    if (u < 0)
      u += M;

    return u;
  }

  inline constexpr static auto frac(int64_t a, int64_t b) {
    return ModInt(a) / ModInt(b);
  }

  inline constexpr auto power(int64_t p) const { return power(val, p); }
  inline constexpr auto inv() const { return inv(val); }

  friend inline constexpr auto operator-(const ModInt &a) {
    return ModInt(-a.val);
  }

  friend inline constexpr auto operator+(int64_t a, const ModInt &b) {
    return ModInt(a) + b;
  }
  friend inline constexpr auto operator-(int64_t a, const ModInt &b) {
    return ModInt(a) - b;
  }
  friend inline constexpr auto operator*(int64_t a, const ModInt &b) {
    return ModInt(a) * b;
  }
  friend inline constexpr auto operator/(int64_t a, const ModInt &b) {
    return ModInt(a) / b;
  }

  friend std::istream &operator>>(std::istream &s, ModInt<M> &a) {
    s >> a.val;
    return s;
  }
  friend std::ostream &operator<<(std::ostream &s, const ModInt<M> &a) {
    s << a.val;
    return s;
  }

  template <int N> inline static auto div() {
    static auto value = inv(N);
    return value;
  }

  explicit operator int32_t() const noexcept { return val; }
  explicit operator int64_t() const noexcept { return val; }
};

using mint = ModInt<998244353>;

/**
 * @title Run length encoding
 * @docs run_length_encoding.md
 */
template <typename Container, typename T = typename Container::value_type>
auto run_length_encoding(const Container &v) {
  std::vector<std::pair<T, int64_t>> ret;

  for (auto &x : v) {
    if (ret.empty())
      ret.push_back({x, 1});
    else if (ret.back().first == x)
      ++ret.back().second;
    else
      ret.push_back({x, 1});
  }

  return ret;
}

mint dp[2][301][150 * 150 + 1];

int main() {
  string S;
  int K;

  while (cin >> S >> K) {
    const int L = count(ALL(S), '1');

    auto r = run_length_encoding(S);
    // dump(r);

    vector<int> a;

    // dump(r);
    if (S[0] == '0')
      a.push_back(0);

    REP(i, r.size()) {
      if (r[i].fst == '1')
        a.push_back(r[i].snd);
      else {
        REP(j, r[i].snd - 1) { a.push_back(0); }
      }
    }

    const int N = a.size();

    dump(a);

    fill_array(dp, 0);

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

    REP(i, N) {
      int cur = i % 2;
      int nxt = (i + 1) % 2;

      FORE(j, 0, L) {
        FORE(k, 0, 150 * 150) {
          FORE(d, max<int>(-j, -a[i]), min<int>(L - j, L - a[i])) {
            if (k + max<int>(0, d) <= 150 * 150) {
              dp[nxt][j + d][k + max<int>(0, d)] += dp[cur][j][k];
            }
          }
        }
      }

      FORE(j, 0, L) {
        FORE(k, 0, 150 * 150) { dp[cur][j][k] = 0; }
      }
    }

    mint ans = 0;

    FORE(i, 0, min(K, 150 * 150)) { ans += dp[N % 2][0][i]; }

    print(ans);
  }

  return 0;
}