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

#include <bits/stdc++.h>
#define rep(i, a, b) for (int i = a; i < b; i++)
#define rrep(i, a, b) for (int i = a; i >= b; i--)
#define fore(i, a) for (auto &i : a)
#define all(x) (x).begin(), (x).end()
// #pragma GCC optimize ("-O3")
using namespace std;
void _main();
int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  _main();
}
typedef long long ll;
const int inf = INT_MAX / 2;
const ll infl = 1LL << 60;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}
//---------------------------------------------------------------------------------------------------
template <int MOD> struct ModInt {
  static const int Mod = MOD;
  unsigned x;
  ModInt() : x(0) {}
  ModInt(signed sig) { x = sig < 0 ? sig % MOD + MOD : sig % MOD; }
  ModInt(signed long long sig) { x = sig < 0 ? sig % MOD + MOD : sig % MOD; }
  int get() const { return (int)x; }
  ModInt &operator+=(ModInt that) {
    if ((x += that.x) >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator-=(ModInt that) {
    if ((x += MOD - that.x) >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator*=(ModInt that) {
    x = (unsigned long long)x * that.x % MOD;
    return *this;
  }
  ModInt &operator/=(ModInt that) { return *this *= that.inverse(); }
  ModInt operator+(ModInt that) const { return ModInt(*this) += that; }
  ModInt operator-(ModInt that) const { return ModInt(*this) -= that; }
  ModInt operator*(ModInt that) const { return ModInt(*this) *= that; }
  ModInt operator/(ModInt that) const { return ModInt(*this) /= that; }
  ModInt inverse() const {
    long long a = x, b = MOD, u = 1, v = 0;
    while (b) {
      long long t = a / b;
      a -= t * b;
      std::swap(a, b);
      u -= t * v;
      std::swap(u, v);
    }
    return ModInt(u);
  }
  bool operator==(ModInt that) const { return x == that.x; }
  bool operator!=(ModInt that) const { return x != that.x; }
  ModInt operator-() const {
    ModInt t;
    t.x = x == 0 ? 0 : Mod - x;
    return t;
  }
};
template <int MOD> ostream &operator<<(ostream &st, const ModInt<MOD> a) {
  st << a.get();
  return st;
};
template <int MOD> ModInt<MOD> operator^(ModInt<MOD> a, unsigned long long k) {
  ModInt<MOD> r = 1;
  while (k) {
    if (k & 1)
      r *= a;
    a *= a;
    k >>= 1;
  }
  return r;
}
typedef ModInt<998244353> mint;
vector<string> split(string str, string delim) { // strをdelimそれぞれで分割
  vector<string> items;
  size_t dlm_idx;
  if (str.npos == (dlm_idx = str.find_first_of(delim))) {
    items.push_back(str.substr(0, dlm_idx));
  }
  while (str.npos != (dlm_idx = str.find_first_of(delim))) {
    if (str.npos == str.find_first_not_of(delim)) {
      break;
    }
    items.push_back(str.substr(0, dlm_idx));
    dlm_idx++;
    str = str.erase(0, dlm_idx);
    if (str.npos == str.find_first_of(delim) && "" != str) {
      items.push_back(str);
      break;
    }
  }
  return items;
}
/*---------------------------------------------------------------------------------------------------
　　　　　　　　　　　 ∧＿∧
　　　　　 ∧＿∧ 　（´<_｀ ）　 Welcome to My Coding Space!
　　　　 （ ´_ゝ`）　/　 ⌒i     @hamayanhamayan0
　　　　／　　　＼　 　  |　|
　　　 /　　 /￣￣￣￣/　　|
　 ＿_(__ﾆつ/　    ＿/ .| .|＿＿＿＿
　 　　　＼/＿＿＿＿/　（u　⊃
---------------------------------------------------------------------------------------------------*/

string S;
int K;
int N;
int A[303];
mint dp[603][603];
mint imos1[603][603], imos2[603][603];
//---------------------------------------------------------------------------------------------------
void _main() {
  cin >> S >> K;
  N = S.length();
  chmin(K, N);

  auto v = split(S, "0");
  int M = v.size();
  rep(i, 0, M) A[i] = v[i].length();

  dp[0][0] = 1;
  rep(i, 0, M) {
    rep(need, 0, N) rep(k, 0, K + 1) imos1[need][k] = 0;
    rep(need, 0, N) rep(k, 0, K + 1) imos2[need][k] = 0;

    /*
               rep(a2, 0, A[i] + 1) if (0 <= need - (A[i] - a2)) dp[ne][need -
       (A[i] - a2)][k + (A[i] - a2)] += dp[cu][need][k]; rep(a2, A[i] + 1, N +
       1) dp[ne][need + (a2 - A[i])][k] += dp[cu][need][k];

    */

    rep(need, 0, N) rep(k, 0, K + 1) {
      int L = min(need, A[i]);
      int R = 0;

      imos1[need - L][k + L] += dp[need][k];
      if (k)
        imos1[need + 1][k - 1] -= dp[need][k];

      imos2[need + 1][k] += dp[need][k];
    }

    rep(need, 1, N) rrep(k, K - 1, 0) imos1[need][k] += imos1[need - 1][k + 1];
    rep(need, 1, N) rep(k, 0, K + 1) imos2[need][k] += imos2[need - 1][k];
    rep(need, 0, N) rep(k, 0, K + 1) dp[need][k] =
        imos1[need][k] + imos2[need][k];
  }

  mint ans = 0;
  rep(k, 0, K + 1) ans += dp[0][k];
  cout << ans << endl;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, a, b) for (int i = a; i < b; i++)
#define rrep(i, a, b) for (int i = a; i >= b; i--)
#define fore(i, a) for (auto &i : a)
#define all(x) (x).begin(), (x).end()
// #pragma GCC optimize ("-O3")
using namespace std;
void _main();
int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  _main();
}
typedef long long ll;
const int inf = INT_MAX / 2;
const ll infl = 1LL << 60;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}
//---------------------------------------------------------------------------------------------------
template <int MOD> struct ModInt {
  static const int Mod = MOD;
  unsigned x;
  ModInt() : x(0) {}
  ModInt(signed sig) { x = sig < 0 ? sig % MOD + MOD : sig % MOD; }
  ModInt(signed long long sig) { x = sig < 0 ? sig % MOD + MOD : sig % MOD; }
  int get() const { return (int)x; }
  ModInt &operator+=(ModInt that) {
    if ((x += that.x) >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator-=(ModInt that) {
    if ((x += MOD - that.x) >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator*=(ModInt that) {
    x = (unsigned long long)x * that.x % MOD;
    return *this;
  }
  ModInt &operator/=(ModInt that) { return *this *= that.inverse(); }
  ModInt operator+(ModInt that) const { return ModInt(*this) += that; }
  ModInt operator-(ModInt that) const { return ModInt(*this) -= that; }
  ModInt operator*(ModInt that) const { return ModInt(*this) *= that; }
  ModInt operator/(ModInt that) const { return ModInt(*this) /= that; }
  ModInt inverse() const {
    long long a = x, b = MOD, u = 1, v = 0;
    while (b) {
      long long t = a / b;
      a -= t * b;
      std::swap(a, b);
      u -= t * v;
      std::swap(u, v);
    }
    return ModInt(u);
  }
  bool operator==(ModInt that) const { return x == that.x; }
  bool operator!=(ModInt that) const { return x != that.x; }
  ModInt operator-() const {
    ModInt t;
    t.x = x == 0 ? 0 : Mod - x;
    return t;
  }
};
template <int MOD> ostream &operator<<(ostream &st, const ModInt<MOD> a) {
  st << a.get();
  return st;
};
template <int MOD> ModInt<MOD> operator^(ModInt<MOD> a, unsigned long long k) {
  ModInt<MOD> r = 1;
  while (k) {
    if (k & 1)
      r *= a;
    a *= a;
    k >>= 1;
  }
  return r;
}
typedef ModInt<998244353> mint;
vector<string> split(string str, string delim) { // strをdelimそれぞれで分割
  vector<string> items;
  size_t dlm_idx;
  if (str.npos == (dlm_idx = str.find_first_of(delim))) {
    items.push_back(str.substr(0, dlm_idx));
  }
  while (str.npos != (dlm_idx = str.find_first_of(delim))) {
    if (str.npos == str.find_first_not_of(delim)) {
      break;
    }
    items.push_back(str.substr(0, dlm_idx));
    dlm_idx++;
    str = str.erase(0, dlm_idx);
    if (str.npos == str.find_first_of(delim) && "" != str) {
      items.push_back(str);
      break;
    }
  }
  return items;
}
/*---------------------------------------------------------------------------------------------------
　　　　　　　　　　　 ∧＿∧
　　　　　 ∧＿∧ 　（´<_｀ ）　 Welcome to My Coding Space!
　　　　 （ ´_ゝ`）　/　 ⌒i     @hamayanhamayan0
　　　　／　　　＼　 　  |　|
　　　 /　　 /￣￣￣￣/　　|
　 ＿_(__ﾆつ/　    ＿/ .| .|＿＿＿＿
　 　　　＼/＿＿＿＿/　（u　⊃
---------------------------------------------------------------------------------------------------*/

string S;
int K;
int N;
int A[303];
mint dp[2][603][603];
//---------------------------------------------------------------------------------------------------
void _main() {
  cin >> S >> K;
  N = S.length();
  chmin(K, N);

  auto v = split(S, "0");
  int M = v.size();
  rep(i, 0, M) A[i] = v[i].length();

  dp[0][0][0] = 1;
  rep(i, 0, M) {
    int cu = i % 2;
    int ne = 1 - cu;
    rep(need, 0, N) rep(k, 0, K + 1) dp[ne][need][k] = 0;
    rep(need, 0, N) rep(k, 0, K + 1) {
      rep(a2, 0, A[i] + 1) if (0 <= need - (A[i] - a2))
          dp[ne][need - (A[i] - a2)][k + (A[i] - a2)] += dp[cu][need][k];
      rep(a2, A[i] + 1, N + 1) dp[ne][need + (a2 - A[i])][k] += dp[cu][need][k];
    }
  }

  mint ans = 0;
  rep(k, 0, K + 1) ans += dp[M % 2][0][k];
  cout << ans << endl;
}