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

/*⠄⢰⣧⣼⣯⠄⣸⣠⣶⣶⣦⣾⠄⠄⠄⠄⡀⠄⢀⣿⣿⠄⠄⠄⢸⡇⠄⠄
⠄⠄⠄⣾⣿⠿⠿⠶⠿⢿⣿⣿⣿⣿⣦⣤⣄⢀⡅⢠⣾⣛⡉⠄⠄⠄⠸⢀⣿⠄
⠄⠄⢀⡋⣡⣴⣶⣶⡀⠄⠄⠙⢿⣿⣿⣿⣿⣿⣴⣿⣿⣿⢃⣤⣄⣀⣥⣿⣿⠄
⠄⠄⢸⣇⠻⣿⣿⣿⣧⣀⢀⣠⡌⢻⣿⣿⣿⣿⣿⣿⣿⣿⣿⠿⠿⠿⣿⣿⣿⠄
⠄⢀⢸⣿⣷⣤⣤⣤⣬⣙⣛⢿⣿⣿⣿⣿⣿⣿⡿⣿⣿⡍⠄⠄⢀⣤⣄⠉⠋⣰
⠄⣼⣖⣿⣿⣿⣿⣿⣿⣿⣿⣿⢿⣿⣿⣿⣿⣿⢇⣿⣿⡷⠶⠶⢿⣿⣿⠇⢀⣤
⠘⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣽⣿⣿⣿⡇⣿⣿⣿⣿⣿⣿⣷⣶⣥⣴⣿⡗
⢀⠈⢿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡟⠄
⢸⣿⣦⣌⣛⣻⣿⣿⣧⠙⠛⠛⡭⠅⠒⠦⠭⣭⡻⣿⣿⣿⣿⣿⣿⣿⣿⡿⠃⠄
⠘⣿⣿⣿⣿⣿⣿⣿⣿⡆⠄⠄⠄⠄⠄⠄⠄⠄⠹⠈⢋⣽⣿⣿⣿⣿⣵⣾⠃⠄
⠄⠘⣿⣿⣿⣿⣿⣿⣿⣿⠄⣴⣿⣶⣄⠄⣴⣶⠄⢀⣾⣿⣿⣿⣿⣿⣿⠃⠄⠄
⠄⠄⠈⠻⣿⣿⣿⣿⣿⣿⡄⢻⣿⣿⣿⠄⣿⣿⡀⣾⣿⣿⣿⣿⣛⠛⠁⠄⠄⠄
⠄⠄⠄⠄⠈⠛⢿⣿⣿⣿⠁⠞⢿⣿⣿⡄⢿⣿⡇⣸⣿⣿⠿⠛⠁⠄⠄⠄⠄⠄
⠄⠄⠄⠄⠄⠄⠄⠉⠻⣿⣿⣾⣦⡙⠻⣷⣾⣿⠃⠿⠋⠁⠄⠄⠄⠄⠄⢀⣠⣴
⣿⣿⣿⣶⣶⣮⣥⣒⠲⢮⣝⡿⣿⣿⡆⣿⡿⠃⠄⠄⠄⠄⠄⠄⠄⣠⣴⣿*/

#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/rope>
using namespace std;
using namespace __gnu_pbds;
using namespace __gnu_cxx;
#define ll long long
#define ii pair<ll, ll>
#define iii pair<ii, ll>
#define fi first
#define se second
#define endl '\\n'
#define debug(x) cout << #x << " is " << x << endl;

#define rep(x, start, end)                                                     \\
  for (auto x = (start) - ((start) > (end)); x != (end) - ((start) > (end));   \\
       ((start) < (end) ? x++ : x--))
#define all(x) (x).begin(), (x).end()
#define sz(x) (int)(x).size()

ll MAX(ll a) { return a; }
ll MIN(ll a) { return a; }
template <typename... Args> ll MAX(ll a, Args... args) {
  return max(a, MAX(args...));
}
template <typename... Args> ll MIN(ll a, Args... args) {
  return min(a, MIN(args...));
}

#define indexed_set                                                            \\
  tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update>

mt19937 rng(chrono::system_clock::now().time_since_epoch().count());

const int MOD = 998244353;

string s;
int k;

vector<int> nums;

ll memo[305][305][305];
ll dp(int i, int used, int left) {
  if (i == sz(nums))
    return used == 0;
  if (memo[i][used][left] != -1)
    return memo[i][used][left];

  ll res = 0;
  rep(x, 1, 305) {
    if (305 <= used + x || left - x < 0)
      break;
    res += dp(i + 1, used + x, left - x);
  }
  rep(x, 0, nums[i] + 1) {
    if (used - x < 0)
      break;
    res += dp(i + 1, used - x, left);
  }

  // cout<<i<<" "<<used<<" "<<left<<" "<<res<<endl;
  return memo[i][used][left] = res % MOD;
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);

  memset(memo, -1, sizeof(memo));

  cin >> s >> k;
  s += '0';

  int curr = 0;
  for (auto &it : s) {
    if (it == '0') {
      nums.push_back(curr);
      curr = 0;
    } else {
      curr++;
    }
  }

  // for (auto &it:nums) cout<<it<<" ";cout<<endl;

  cout << dp(0, 0, min(k, 304)) << endl;
}


Error: TLE

Fixed submission:

/*⠄⢰⣧⣼⣯⠄⣸⣠⣶⣶⣦⣾⠄⠄⠄⠄⡀⠄⢀⣿⣿⠄⠄⠄⢸⡇⠄⠄
⠄⠄⠄⣾⣿⠿⠿⠶⠿⢿⣿⣿⣿⣿⣦⣤⣄⢀⡅⢠⣾⣛⡉⠄⠄⠄⠸⢀⣿⠄
⠄⠄⢀⡋⣡⣴⣶⣶⡀⠄⠄⠙⢿⣿⣿⣿⣿⣿⣴⣿⣿⣿⢃⣤⣄⣀⣥⣿⣿⠄
⠄⠄⢸⣇⠻⣿⣿⣿⣧⣀⢀⣠⡌⢻⣿⣿⣿⣿⣿⣿⣿⣿⣿⠿⠿⠿⣿⣿⣿⠄
⠄⢀⢸⣿⣷⣤⣤⣤⣬⣙⣛⢿⣿⣿⣿⣿⣿⣿⡿⣿⣿⡍⠄⠄⢀⣤⣄⠉⠋⣰
⠄⣼⣖⣿⣿⣿⣿⣿⣿⣿⣿⣿⢿⣿⣿⣿⣿⣿⢇⣿⣿⡷⠶⠶⢿⣿⣿⠇⢀⣤
⠘⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣽⣿⣿⣿⡇⣿⣿⣿⣿⣿⣿⣷⣶⣥⣴⣿⡗
⢀⠈⢿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡟⠄
⢸⣿⣦⣌⣛⣻⣿⣿⣧⠙⠛⠛⡭⠅⠒⠦⠭⣭⡻⣿⣿⣿⣿⣿⣿⣿⣿⡿⠃⠄
⠘⣿⣿⣿⣿⣿⣿⣿⣿⡆⠄⠄⠄⠄⠄⠄⠄⠄⠹⠈⢋⣽⣿⣿⣿⣿⣵⣾⠃⠄
⠄⠘⣿⣿⣿⣿⣿⣿⣿⣿⠄⣴⣿⣶⣄⠄⣴⣶⠄⢀⣾⣿⣿⣿⣿⣿⣿⠃⠄⠄
⠄⠄⠈⠻⣿⣿⣿⣿⣿⣿⡄⢻⣿⣿⣿⠄⣿⣿⡀⣾⣿⣿⣿⣿⣛⠛⠁⠄⠄⠄
⠄⠄⠄⠄⠈⠛⢿⣿⣿⣿⠁⠞⢿⣿⣿⡄⢿⣿⡇⣸⣿⣿⠿⠛⠁⠄⠄⠄⠄⠄
⠄⠄⠄⠄⠄⠄⠄⠉⠻⣿⣿⣾⣦⡙⠻⣷⣾⣿⠃⠿⠋⠁⠄⠄⠄⠄⠄⢀⣠⣴
⣿⣿⣿⣶⣶⣮⣥⣒⠲⢮⣝⡿⣿⣿⡆⣿⡿⠃⠄⠄⠄⠄⠄⠄⠄⣠⣴⣿*/

#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/rope>
using namespace std;
using namespace __gnu_pbds;
using namespace __gnu_cxx;
#define ll long long
#define ii pair<ll, ll>
#define iii pair<ii, ll>
#define fi first
#define se second
#define endl '\\n'
#define debug(x) cout << #x << " is " << x << endl;

#define rep(x, start, end)                                                     \\
  for (auto x = (start) - ((start) > (end)); x != (end) - ((start) > (end));   \\
       ((start) < (end) ? x++ : x--))
#define all(x) (x).begin(), (x).end()
#define sz(x) (int)(x).size()

ll MAX(ll a) { return a; }
ll MIN(ll a) { return a; }
template <typename... Args> ll MAX(ll a, Args... args) {
  return max(a, MAX(args...));
}
template <typename... Args> ll MIN(ll a, Args... args) {
  return min(a, MIN(args...));
}

#define indexed_set                                                            \\
  tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update>

mt19937 rng(chrono::system_clock::now().time_since_epoch().count());

const int MOD = 998244353;

string s;
int k;

vector<int> nums;

ll fix(ll i) {
  i %= MOD;
  if (i < 0)
    i += MOD;
  return i;
}

ll memo[305][305][305];
ll dp(int i, int used, int left) {
  if (i == sz(nums))
    return used == 0;
  if (memo[i][used][left] != -1)
    return memo[i][used][left];

  int res = 0;
  rep(x, 1, 305) {
    if (305 <= used + x || left - x < 0)
      break;
    res = fix(res + dp(i + 1, used + x, left - x));
  }
  rep(x, 0, nums[i] + 1) {
    if (used - x < 0)
      break;
    res = fix(res + dp(i + 1, used - x, left));
  }

  // cout<<i<<" "<<used<<" "<<left<<" "<<res<<endl;
  return memo[i][used][left] = res;
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);

  memset(memo, -1, sizeof(memo));

  cin >> s >> k;
  s += '0';

  int curr = 0;
  for (auto &it : s) {
    if (it == '0') {
      nums.push_back(curr);
      curr = 0;
    } else {
      curr++;
    }
  }

  // for (auto &it:nums) cout<<it<<" ";cout<<endl;

  cout << dp(0, 0, min(k, 304)) << endl;
}