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

#ifdef LOCAL
// #define _GLIBCXX_DEBUG
#endif
// #pragma GCC target("avx512f,avx512dq,avx512cd,avx512bw,avx512vl")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#include <bits/stdc++.h>

using namespace std;

typedef int ll;
typedef pair<ll, ll> P;
typedef pair<int, int> Pi;
typedef vector<ll> Vec;
typedef vector<int> Vi;
typedef vector<string> Vs;
typedef vector<char> Vc;
typedef vector<P> VP;
typedef vector<vector<ll>> VV;
typedef vector<vector<int>> VVi;
typedef vector<vector<char>> VVc;
typedef vector<vector<vector<ll>>> VVV;
typedef vector<vector<vector<vector<ll>>>> VVVV;

#define REP(i, a, b) for (ll i = (a); i < (b); i++)
#define PER(i, a, b) for (ll i = (a); i >= (b); i--)
#define rep(i, n) REP(i, 0, n)
#define per(i, n) PER(i, n, 0)
const ll INF = 1e18 + 18;
const ll MOD = 998244353;
#define Yes(n) cout << ((n) ? "Yes" : "No") << endl;
#define YES(n) cout << ((n) ? "YES" : "NO") << endl;
#define ALL(v) v.begin(), v.end()
#define rALL(v) v.rbegin(), v.rend()
#define pb(x) push_back(x)
#define mp(a, b) make_pair(a, b)
#define Each(a, b) for (auto &a : b)
#define rEach(i, mp) for (auto i = mp.rbegin(); i != mp.rend(); ++i)
#ifdef LOCAL
#define dbg(x_) cerr << #x_ << ":" << x_ << endl;
#define dbgmap(mp)                                                             \\
  cerr << #mp << ":" << endl;                                                  \\
  for (auto i = mp.begin(); i != mp.end(); ++i) {                              \\
    cerr << i->first << ":" << i->second << endl;                              \\
  }
#define dbgarr(n, m, arr)                                                      \\
  rep(i, n) {                                                                  \\
    rep(j, m) { cerr << arr[i][j] << " "; }                                    \\
    cerr << endl;                                                              \\
  }
#define dbgdp(n, arr)                                                          \\
  rep(i, n) { cerr << arr[i] << " "; }                                         \\
  cerr << endl;
#define dbgmint(n, arr)                                                        \\
  rep(i, n) { cerr << arr[i].x << " "; }                                       \\
  cerr << endl;
#define dbgarrmint(n, m, arr)                                                  \\
  rep(i, n) {                                                                  \\
    rep(j, m) { cerr << arr[i][j].x << " "; }                                  \\
    cerr << endl;                                                              \\
  }
#else
#define dbg(...)
#define dbgmap(...)
#define dbgarr(...)
#define dbgdp(...)
#define dbgmint(...)
#define dbgarrmint(...)
#endif
#define out(a) cout << a << endl
#define out2(a, b) cout << a << " " << b << endl
#define vout(v)                                                                \\
  rep(i, v.size()) { cout << v[i] << " "; }                                    \\
  cout << endl
#define Uniq(v) v.erase(unique(v.begin(), v.end()), v.end())
#define fi first
#define se 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;
}

template <typename T1, typename T2>
ostream &operator<<(ostream &s, const pair<T1, T2> &p) {
  return s << "(" << p.first << ", " << p.second << ")";
}

template <typename T> istream &operator>>(istream &i, vector<T> &v) {
  rep(j, v.size()) i >> v[j];
  return i;
}

// vector
template <typename T> ostream &operator<<(ostream &s, const vector<T> &v) {
  int len = v.size();
  for (int i = 0; i < len; ++i) {
    s << v[i];
    if (i < len - 1)
      s << "	";
  }
  return s;
}

// 2 dimentional vector
template <typename T>
ostream &operator<<(ostream &s, const vector<vector<T>> &vv) {
  int len = vv.size();
  for (int i = 0; i < len; ++i) {
    s << vv[i] << endl;
  }
  return s;
}

// ランレングス
vector<pair<char, ll>> run_length_encoding(const string &s) {
  vector<pair<char, ll>> ans;
  ll num = 1;
  rep(i, s.length() - 1) {
    if (s[i] == s[i + 1]) {
      num++;
    } else {
      ans.emplace_back(mp(s[i], num));
      num = 1;
    }
  }
  ans.emplace_back(mp(s.back(), num));
  return ans;
}

void simulate() {
  // string s = "01100110";
  string s = "01010101";
  ll n = s.length();
  ll k = 2;
  map<string, ll> mp;
  queue<pair<string, ll>> q;
  q.emplace(s, 0);
  while (!q.empty()) {
    pair<string, ll> p = q.front();
    string &t = p.fi;
    mp[t]++;
    // if(mp[t]>1)continue;
    q.pop();
    if (p.se >= k)
      continue;
    rep(j, n) rep(i, j) {
      if (t[i] == '0' && t[j] == '1') {
        string x;
        rep(id, n) {
          if (id == j)
            continue;
          if (id == i) {
            x += '1';
            x += '0';
          } else {
            x += t[id];
          }
        }
        q.emplace(x, p.se + 1);
      }
    }
  }
  dbgmap(mp);
  dbg(mp.size());
  Each(x, mp) {
    string t = x.fi;
    vector<pair<char, ll>> rl = run_length_encoding(t);
    Vec v(n);
    ll id = 0;
    Each(p, rl) {
      if (p.fi == '0') {
        id += p.se;
      } else {
        v[id] = p.se;
      }
    }
    dbg(v);
  }
  exit(0);
}
// mint
struct mint {
  ll x; // typedef long long ll;
  mint(ll x = 0) : x((x % MOD + MOD) % MOD) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint a) {
    if ((x += a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += MOD - a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= MOD;
    return *this;
  }
  mint operator+(const mint a) const { return mint(*this) += a; }
  mint operator-(const mint a) const { return mint(*this) -= a; }
  mint operator*(const mint a) const { return mint(*this) *= a; }
  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime MOD
  mint inv() const { return pow(MOD - 2); }
  mint &operator/=(const mint a) { return *this *= a.inv(); }
  mint operator/(const mint a) const { return mint(*this) /= a; }
};
istream &operator>>(istream &is, const mint &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const mint &a) { return os << a.x; }
typedef vector<mint> Vmint;

struct combination {
  vector<mint> fact, ifact;
  combination(int n) : fact(n + 1), ifact(n + 1) {
    assert(n < MOD);
    fact[0] = 1;
    for (int i = 1; i <= n; ++i)
      fact[i] = fact[i - 1] * i;
    ifact[n] = fact[n].inv();
    for (int i = n; i >= 1; --i)
      ifact[i - 1] = ifact[i] * i;
  }
  mint operator()(int n, int k) {
    if (k < 0 || k > n)
      return 0;
    return fact[n] * ifact[k] * ifact[n - k];
  }
};

int solve() {
  string s;
  ll k;
  cin >> s >> k;
  vector<pair<char, ll>> rl = run_length_encoding(s);
  ll n = s.length();
  k = min(n, k);
  dbg(rl);
  Vec v(n);
  ll id = 0;
  Each(p, rl) {
    if (p.fi == '0') {
      id += p.se;
    } else {
      v[id] = p.se;
    }
  }

  reverse(ALL(v));
  dbg(v);

  vector<vector<Vmint>> dp(n + 1, vector<Vmint>(k + 1, Vmint(k + 1, 0)));
  dp[0][0][0] = 1;
  // i番目まででsub個引いてadd足した場合。
  rep(i, n) {
    rep(sub, k + 1) {
      rep(add, sub + 1) {
        // sub
        rep(cnt, min(v[i] + 1, k - sub + 1)) {
          //   if(sub+cnt>k)continue;
          dp[i + 1][sub + cnt][add] += dp[i][sub][add];
        }
        // add
        REP(cnt, 1, min(k + 1, k - add + 1)) {
          //  if(add+cnt>k)continue;
          dp[i + 1][sub][add + cnt] += dp[i][sub][add];
        }
      }
    }
  }
  mint ans = 0;
  // dbg(dp);
  rep(i, k + 1) { ans += dp[n][i][i]; }
  out(ans);
  return 0;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << std::setprecision(10);
  // simulate();
  solve();
}


Error: TLE

Fixed submission:

#ifdef LOCAL
// #define _GLIBCXX_DEBUG
#endif
// #pragma GCC target("avx512f,avx512dq,avx512cd,avx512bw,avx512vl")
// #pragma GCC optimize("O3")
// #pragma GCC optimize("unroll-loops")
#include <bits/stdc++.h>

using namespace std;

typedef int ll;
typedef pair<ll, ll> P;
typedef pair<int, int> Pi;
typedef vector<ll> Vec;
typedef vector<int> Vi;
typedef vector<string> Vs;
typedef vector<char> Vc;
typedef vector<P> VP;
typedef vector<vector<ll>> VV;
typedef vector<vector<int>> VVi;
typedef vector<vector<char>> VVc;
typedef vector<vector<vector<ll>>> VVV;
typedef vector<vector<vector<vector<ll>>>> VVVV;

#define REP(i, a, b) for (ll i = (a); i < (b); i++)
#define PER(i, a, b) for (ll i = (a); i >= (b); i--)
#define rep(i, n) REP(i, 0, n)
#define per(i, n) PER(i, n, 0)
const ll INF = 1e18 + 18;
const ll MOD = 998244353;
#define Yes(n) cout << ((n) ? "Yes" : "No") << endl;
#define YES(n) cout << ((n) ? "YES" : "NO") << endl;
#define ALL(v) v.begin(), v.end()
#define rALL(v) v.rbegin(), v.rend()
#define pb(x) push_back(x)
#define mp(a, b) make_pair(a, b)
#define Each(a, b) for (auto &a : b)
#define rEach(i, mp) for (auto i = mp.rbegin(); i != mp.rend(); ++i)
#ifdef LOCAL
#define dbg(x_) cerr << #x_ << ":" << x_ << endl;
#define dbgmap(mp)                                                             \\
  cerr << #mp << ":" << endl;                                                  \\
  for (auto i = mp.begin(); i != mp.end(); ++i) {                              \\
    cerr << i->first << ":" << i->second << endl;                              \\
  }
#define dbgarr(n, m, arr)                                                      \\
  rep(i, n) {                                                                  \\
    rep(j, m) { cerr << arr[i][j] << " "; }                                    \\
    cerr << endl;                                                              \\
  }
#define dbgdp(n, arr)                                                          \\
  rep(i, n) { cerr << arr[i] << " "; }                                         \\
  cerr << endl;
#define dbgmint(n, arr)                                                        \\
  rep(i, n) { cerr << arr[i].x << " "; }                                       \\
  cerr << endl;
#define dbgarrmint(n, m, arr)                                                  \\
  rep(i, n) {                                                                  \\
    rep(j, m) { cerr << arr[i][j].x << " "; }                                  \\
    cerr << endl;                                                              \\
  }
#else
#define dbg(...)
#define dbgmap(...)
#define dbgarr(...)
#define dbgdp(...)
#define dbgmint(...)
#define dbgarrmint(...)
#endif
#define out(a) cout << a << endl
#define out2(a, b) cout << a << " " << b << endl
#define vout(v)                                                                \\
  rep(i, v.size()) { cout << v[i] << " "; }                                    \\
  cout << endl
#define Uniq(v) v.erase(unique(v.begin(), v.end()), v.end())
#define fi first
#define se 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;
}

template <typename T1, typename T2>
ostream &operator<<(ostream &s, const pair<T1, T2> &p) {
  return s << "(" << p.first << ", " << p.second << ")";
}

template <typename T> istream &operator>>(istream &i, vector<T> &v) {
  rep(j, v.size()) i >> v[j];
  return i;
}

// vector
template <typename T> ostream &operator<<(ostream &s, const vector<T> &v) {
  int len = v.size();
  for (int i = 0; i < len; ++i) {
    s << v[i];
    if (i < len - 1)
      s << "	";
  }
  return s;
}

// 2 dimentional vector
template <typename T>
ostream &operator<<(ostream &s, const vector<vector<T>> &vv) {
  int len = vv.size();
  for (int i = 0; i < len; ++i) {
    s << vv[i] << endl;
  }
  return s;
}

// ランレングス
vector<pair<char, ll>> run_length_encoding(const string &s) {
  vector<pair<char, ll>> ans;
  ll num = 1;
  rep(i, s.length() - 1) {
    if (s[i] == s[i + 1]) {
      num++;
    } else {
      ans.emplace_back(mp(s[i], num));
      num = 1;
    }
  }
  ans.emplace_back(mp(s.back(), num));
  return ans;
}

void simulate() {
  // string s = "01100110";
  string s = "01010101";
  ll n = s.length();
  ll k = 2;
  map<string, ll> mp;
  queue<pair<string, ll>> q;
  q.emplace(s, 0);
  while (!q.empty()) {
    pair<string, ll> p = q.front();
    string &t = p.fi;
    mp[t]++;
    // if(mp[t]>1)continue;
    q.pop();
    if (p.se >= k)
      continue;
    rep(j, n) rep(i, j) {
      if (t[i] == '0' && t[j] == '1') {
        string x;
        rep(id, n) {
          if (id == j)
            continue;
          if (id == i) {
            x += '1';
            x += '0';
          } else {
            x += t[id];
          }
        }
        q.emplace(x, p.se + 1);
      }
    }
  }
  dbgmap(mp);
  dbg(mp.size());
  Each(x, mp) {
    string t = x.fi;
    vector<pair<char, ll>> rl = run_length_encoding(t);
    Vec v(n);
    ll id = 0;
    Each(p, rl) {
      if (p.fi == '0') {
        id += p.se;
      } else {
        v[id] = p.se;
      }
    }
    dbg(v);
  }
  exit(0);
}
// mint
struct mint {
  ll x; // typedef long long ll;
  mint(ll x = 0) : x((x % MOD + MOD) % MOD) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint a) {
    if ((x += a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += MOD - a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= MOD;
    return *this;
  }
  mint operator+(const mint a) const { return mint(*this) += a; }
  mint operator-(const mint a) const { return mint(*this) -= a; }
  mint operator*(const mint a) const { return mint(*this) *= a; }
  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime MOD
  mint inv() const { return pow(MOD - 2); }
  mint &operator/=(const mint a) { return *this *= a.inv(); }
  mint operator/(const mint a) const { return mint(*this) /= a; }
};
istream &operator>>(istream &is, const mint &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const mint &a) { return os << a.x; }
typedef vector<mint> Vmint;

struct combination {
  vector<mint> fact, ifact;
  combination(int n) : fact(n + 1), ifact(n + 1) {
    assert(n < MOD);
    fact[0] = 1;
    for (int i = 1; i <= n; ++i)
      fact[i] = fact[i - 1] * i;
    ifact[n] = fact[n].inv();
    for (int i = n; i >= 1; --i)
      ifact[i - 1] = ifact[i] * i;
  }
  mint operator()(int n, int k) {
    if (k < 0 || k > n)
      return 0;
    return fact[n] * ifact[k] * ifact[n - k];
  }
};

int solve() {
  string s;
  ll k;
  cin >> s >> k;
  vector<pair<char, ll>> rl = run_length_encoding(s);
  ll n = s.length();
  k = min(n, k);
  dbg(rl);
  Vec v(n);
  ll id = 0;
  Each(p, rl) {
    if (p.fi == '0') {
      id += p.se;
    } else {
      v[id] = p.se;
    }
  }

  reverse(ALL(v));
  dbg(v);

  vector<vector<Vmint>> dp(n + 1, vector<Vmint>(k + 1, Vmint(k + 1, 0)));
  dp[0][0][0] = 1;
  // i番目まででsub個引いてadd足した場合。
  rep(i, n) {
    rep(sub, k + 1) {
      rep(add, sub + 1) {
        // sub
        rep(cnt, v[i] + 1) {
          if (sub + cnt > k)
            continue;
          dp[i + 1][sub + cnt][add] += dp[i][sub][add];
        }
        // add
        REP(cnt, 1, k + 1) {
          if (add + cnt > k)
            continue;
          dp[i + 1][sub][add + cnt] += dp[i][sub][add];
        }
      }
    }
  }
  mint ans = 0;
  // dbg(dp);
  rep(i, k + 1) { ans += dp[n][i][i]; }
  out(ans);
  return 0;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << std::setprecision(10);
  // simulate();
  solve();
}