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

#include <bits/stdc++.h>
using namespace std;
const int N = 310;
const int mod = 998244353;
typedef vector<int> vi;
#define pb push_back
#define sz(a) int(a.size())

char s[N];
int n, f[N][N][N], g[N][N][N], h[N][N][N], sum[N][2];

int getf(int i, int j, int k) {
  if (j + k > n - i)
    return 0;
  if (i == n)
    return 1;
  if (f[i][j][k] != -1)
    return f[i][j][k];
  for (int a = 0; a < 2; a++)
    for (int b = 0; b < 2; b++)
      for (int x = 0; x < 2; x++)
        for (int y = 0; y < 2; y++) {
          int st[3], tp = 0;
          if (!x)
            st[++tp] = a;
          if (!y)
            st[++tp] = b;
          bool fl = 1;
          for (int l = 1; l <= tp; l++)
            if (s[i + l] != st[l]) {
              fl = 0;
              break;
            }
          if (!fl)
            continue;
          int cnt[2];
          cnt[0] = j;
          cnt[1] = k;
          if (cnt[b]) {
            --cnt[b];
            if (x)
              ++cnt[a];
            if (y)
              ++cnt[b];
            if (getf(i + tp, cnt[0], cnt[1]))
              return f[i][j][k] = 1;
          }
          cnt[0] = j;
          cnt[1] = k;
          if (cnt[a]) {
            --cnt[a];
            if (x)
              ++cnt[a];
            if (y)
              ++cnt[b];
            if (getf(i + tp, cnt[0], cnt[1]))
              return f[i][j][k] = 1;
          }
        }
  return f[i][j][k] = 0;
}

int geth(int i, int j, int k) {
  if (i < 0)
    return 0;
  if (h[i][j][k] != -1)
    return h[i][j][k];
  if (j >= sum[i][0] && k >= sum[i][1] && getf(i, j - sum[i][0], k - sum[i][1]))
    return h[i][j][k] = 1;
  return h[i][j][k] = geth(i - 1, j, k);
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cin >> (s + 1);
  n = strlen(s + 1);
  reverse(s + 1, s + n + 1);
  for (int i = 1; i <= n; i++)
    s[i] -= '0', sum[i][0] = sum[i - 1][0], sum[i][1] = sum[i - 1][1],
                 ++sum[i][s[i]];
  s[n + 1] = 2;
  memset(f, -1, sizeof(f));
  memset(h, -1, sizeof(h));
  int ans = 0;
  g[0][0][0] = 1;
  for (int i = 0; i < n; i++)
    for (int j = 0; i + j < n; j++)
      for (int k = 0; k <= i + j; k++) {
        if (!g[i][j][k])
          continue;
        g[i + 1][j][k + (s[k + 1] == 0)] =
            (g[i + 1][j][k + (s[k + 1] == 0)] + g[i][j][k]) % mod;
        g[i][j + 1][k + (s[k + 1] == 1)] =
            (g[i][j + 1][k + (s[k + 1] == 1)] + g[i][j][k]) % mod;
      }
  for (int i = 0; i <= n; i++)
    for (int j = 0; i + j <= n; j++)
      for (int k = 0; k <= i + j; k++)
        if (geth(k, i, j))
          ans = (ans + g[i][j][k]) % mod;
  cout << ans;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
const int N = 310;
const int mod = 998244353;
typedef vector<int> vi;
#define pb push_back
#define sz(a) int(a.size())

char s[N];
int n, f[N][N][N], g[N][N][N], h[N][N][N], sum[N][2];

int getf(int i, int j, int k) {
  if (j + k > n - i)
    return 0;
  if (i == n)
    return 1;
  if (f[i][j][k] != -1)
    return f[i][j][k];
  for (int a = 0; a < 2; a++)
    for (int b = 0; b < 2; b++)
      for (int x = 0; x < 2; x++)
        for (int y = 0; y < 2; y++) {
          vi vec;
          if (!x)
            vec.pb(a);
          if (!y)
            vec.pb(b);
          bool fl = 1;
          for (int l = 0; l < sz(vec); l++)
            if (s[i + 1 + l] != vec[l]) {
              fl = 0;
              break;
            }
          if (!fl)
            continue;
          int cnt[2];
          cnt[0] = j;
          cnt[1] = k;
          if (cnt[b]) {
            --cnt[b];
            if (x)
              ++cnt[a];
            if (y)
              ++cnt[b];
            if (getf(i + sz(vec), cnt[0], cnt[1]))
              return f[i][j][k] = 1;
          }
          cnt[0] = j;
          cnt[1] = k;
          if (cnt[a]) {
            --cnt[a];
            if (x)
              ++cnt[a];
            if (y)
              ++cnt[b];
            if (getf(i + sz(vec), cnt[0], cnt[1]))
              return f[i][j][k] = 1;
          }
        }
  return f[i][j][k] = 0;
}

int geth(int i, int j, int k) {
  if (i < 0)
    return 0;
  if (h[i][j][k] != -1)
    return h[i][j][k];
  if (j >= sum[i][0] && k >= sum[i][1] && getf(i, j - sum[i][0], k - sum[i][1]))
    return h[i][j][k] = 1;
  return h[i][j][k] = geth(i - 1, j, k);
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cin >> (s + 1);
  n = strlen(s + 1);
  reverse(s + 1, s + n + 1);
  for (int i = 1; i <= n; i++)
    s[i] -= '0', sum[i][0] = sum[i - 1][0], sum[i][1] = sum[i - 1][1],
                 ++sum[i][s[i]];
  s[n + 1] = 2;
  memset(f, -1, sizeof(f));
  memset(h, -1, sizeof(h));
  int ans = 0;
  g[0][0][0] = 1;
  for (int i = 0; i < n; i++)
    for (int j = 0; i + j < n; j++)
      for (int k = 0; k <= i + j; k++) {
        g[i + 1][j][k + (s[k + 1] == 0)] =
            (g[i + 1][j][k + (s[k + 1] == 0)] + g[i][j][k]) % mod;
        g[i][j + 1][k + (s[k + 1] == 1)] =
            (g[i][j + 1][k + (s[k + 1] == 1)] + g[i][j][k]) % mod;
      }
  for (int i = 0; i <= n; i++)
    for (int j = 0; i + j <= n; j++)
      for (int k = 0; k <= i + j; k++)
        if (geth(k, i, j))
          ans = (ans + g[i][j][k]) % mod;
  cout << ans;
  return 0;
}