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

#include <bits/stdc++.h>
using namespace std;

#define IOS                                                                    \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define endl "\\n"

const int N = 301;
const int MOD = 998244353;

int n, k;
string s;
vector<int> v, suf;
int cache[N][N][N];

void add(int &a, int b) {
  a += b;
  if (a >= MOD)
    a -= MOD;
}

int dp(int idx, int cur, int tot) {
  if (tot > k)
    return 0;
  if (idx >= v.size())
    return (cur == 0);
  int &ans = cache[idx][cur][tot];
  if (ans != -1)
    return ans;
  ans = 0;
  int mx = v[idx] + k - tot;
  mx = min(mx, suf[idx]);
  for (int len = max(0, v[idx] - cur); len <= mx; len++)
    add(ans, dp(idx + 1, cur + len - v[idx], tot + max(0, len - v[idx])));
  return ans;
}

int32_t main() {
  IOS;
  memset(cache, -1, sizeof(cache));
  cin >> s >> k;
  while (s.size() && s.back() == '0')
    s.pop_back();
  reverse(s.begin(), s.end());
  while (s.size() && s.back() == '1')
    s.pop_back();
  reverse(s.begin(), s.end());
  int cnt = 0;
  n = s.size();
  k = min(k, n);
  v.push_back(0);
  int prv = 0;
  for (int i = 0; i < n; i++) {
    if (s[i] == '0') {
      if (i > 0 && s[i - 1] == '0')
        v.push_back(0);
      continue;
    }
    int j = i;
    while (j + 1 < n && s[j + 1] == s[i])
      j++;
    v.push_back(j - i + 1);
    cnt += j - i + 1;
    i = j;
  }
  suf = v;
  for (int i = (int)suf.size() - 2; i >= 0; i--)
    suf[i] += suf[i + 1];
  int ans = dp(0, 0, 0);
  cout << ans;
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

#define IOS                                                                    \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define endl "\\n"
#define int long long

const int N = 301;
const int MOD = 998244353;

int n, k;
string s;
vector<int> v, suf;
int cache[N][N][N];

int dp(int idx, int cur, int tot) {
  if (tot > k)
    return 0;
  if (idx >= v.size())
    return (cur == 0);
  int &ans = cache[idx][cur][tot];
  if (ans != -1)
    return ans;
  ans = 0;
  for (int len = max(0LL, v[idx] - cur); len <= v[idx] + k - tot; len++)
    ans += dp(idx + 1, cur + len - v[idx], tot + max(0LL, len - v[idx]));
  ans %= MOD;
  return ans;
}

int32_t main() {
  IOS;
  memset(cache, -1, sizeof(cache));
  cin >> s >> k;
  while (s.size() && s.back() == '0')
    s.pop_back();
  reverse(s.begin(), s.end());
  while (s.size() && s.back() == '1')
    s.pop_back();
  reverse(s.begin(), s.end());
  int cnt = 0;
  n = s.size();
  k = min(k, n);
  v.push_back(0);
  int prv = 0;
  for (int i = 0; i < n; i++) {
    if (s[i] == '0') {
      if (i > 0 && s[i - 1] == '0')
        v.push_back(0);
      continue;
    }
    int j = i;
    while (j + 1 < n && s[j + 1] == s[i])
      j++;
    v.push_back(j - i + 1);
    cnt += j - i + 1;
    i = j;
  }
  suf = v;
  for (int i = (int)suf.size() - 2; i >= 0; i--)
    suf[i] += suf[i + 1];
  int ans = dp(0, 0, 0);
  cout << ans;
  return 0;
}