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

#include <algorithm>
#include <cstring>
#include <set>
#include <stdio.h>
#include <vector>
using namespace std;

const long long mod = 998244353;
char S[303];
int L;

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

pair<int, int> chk[303][303];
bool seen[303][303];

void go(int i, int z, int o, int f) {
  if (seen[L - i + z + o + f][i]) {
    auto &p = chk[L - i + z + o + f][i];
    if (p.first <= z && z + f <= p.second)
      return;
    p.first = min(p.first, z);
    p.second = max(p.second, z + f);
  } else {
    seen[L - i + z + o + f][i] = 1;
    chk[L - i + z + o + f][i] = {z, z + f};
  }
  if (i == L)
    return;

  if (S[i] == '0' && o)
    go(i + 1, z, o - 1, f + 1);
  else if (S[i] == '0' && (z || f))
    go(i + 1, z, o, f);

  if (S[i] == '1' && z)
    go(i + 1, z - 1, o, f + 1);
  else if (S[i] == '1' && (o || f))
    go(i + 1, z, o, f);

  if (i + 1 < L && S[i] == '0' && S[i + 1] == '0')
    go(i + 2, z + 1, o, f);
  if (i + 1 < L && S[i] == '1' && S[i + 1] == '1')
    go(i + 2, z, o + 1, f);
  if (i + 1 < L && S[i] != S[i + 1])
    go(i + 2, z, o, f + 1);
}

long long D[303][303][303], R[303][303][303];
int zc[303], oc[303], pin[303];

int main() {
  scanf("%s", S);
  while (S[L])
    L++;
  for (int i = L - 1; i >= 0; i--) {
    zc[i] = zc[i + 1];
    oc[i] = oc[i + 1];
    if (S[i] == '0')
      zc[i]++;
    if (S[i] == '1')
      oc[i]++;
  }

  go(0, 0, 0, 0);

  D[L][0][0] = 1;
  for (int i = L; i >= 0; i--) {
    for (int j = 0; j <= i; j++)
      for (int k = 0; j + k <= i; k++) {
        R[i][j][k] = D[i][j][k];
        if (j)
          add(R[i][j][k], R[i][j - 1][k]);
        if (k)
          add(R[i][j][k], R[i][j][k - 1]);
      }
    for (int j = 0; j <= i; j++)
      for (int k = 0; j + k <= i; k++) {
        if (i == 0 || S[i - 1] != '0')
          add(D[i][j + 1][k], D[i][j][k]);
        else
          add(D[i - 1][j][k], D[i][j][k]);
        if (i == 0 || S[i - 1] != '1')
          add(D[i][j][k + 1], D[i][j][k]);
        else
          add(D[i - 1][j][k], D[i][j][k]);
      }
  }

  long long ans = 0;
  for (int l = 1; l <= L; l++) {
    memset(pin, -1, sizeof(pin));
    for (int i = 0; i <= L; i++)
      if (seen[l][i]) {
        for (int j = chk[l][i].first; j <= chk[l][i].second; j++)
          pin[j + zc[i]] = i;
      }

    long long ret = 0;
    for (int j = 0; j <= L; j++)
      if (pin[j] != -1) {
        int p = pin[j];
        int z = j - zc[p];
        int o = l - j - oc[p];
        if (o >= 0) {
          add(ans, R[p][z][o]);
          add(ret, R[p][z][o]);
        }
      }
  }

  printf("%lld\\n", ans);

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cstring>
#include <set>
#include <stdio.h>
#include <vector>
using namespace std;

const long long mod = 998244353;
char S[303];
int L;

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

set<pair<int, int>> chk[303][303];

void go(int i, int z, int o, int f) {
  int u = 0;
  if (chk[L - i + z + o + f][i].count({z, z + f}))
    return;
  chk[L - i + z + o + f][i].insert({z, z + f});
  if (i == L)
    return;

  if (S[i] == '0' && o)
    go(i + 1, z, o - 1, f + 1);
  else if (S[i] == '0' && (z || f))
    go(i + 1, z, o, f);

  if (S[i] == '1' && z)
    go(i + 1, z - 1, o, f + 1);
  else if (S[i] == '1' && (o || f))
    go(i + 1, z, o, f);

  if (i + 1 < L && S[i] == '0' && S[i + 1] == '0')
    go(i + 2, z + 1, o, f);
  if (i + 1 < L && S[i] == '1' && S[i + 1] == '1')
    go(i + 2, z, o + 1, f);
  if (i + 1 < L && S[i] != S[i + 1])
    go(i + 2, z, o, f + 1);
}

long long D[303][303][303], R[303][303][303];
int zc[303], oc[303], pin[303];

int main() {
  scanf("%s", S);
  while (S[L])
    L++;
  for (int i = L - 1; i >= 0; i--) {
    zc[i] = zc[i + 1];
    oc[i] = oc[i + 1];
    if (S[i] == '0')
      zc[i]++;
    if (S[i] == '1')
      oc[i]++;
  }

  go(0, 0, 0, 0);

  D[L][0][0] = 1;
  for (int i = L; i >= 0; i--) {
    for (int j = 0; j <= i; j++)
      for (int k = 0; j + k <= i; k++) {
        R[i][j][k] = D[i][j][k];
        if (j)
          add(R[i][j][k], R[i][j - 1][k]);
        if (k)
          add(R[i][j][k], R[i][j][k - 1]);
      }
    for (int j = 0; j <= i; j++)
      for (int k = 0; j + k <= i; k++) {
        if (i == 0 || S[i - 1] != '0')
          add(D[i][j + 1][k], D[i][j][k]);
        else
          add(D[i - 1][j][k], D[i][j][k]);
        if (i == 0 || S[i - 1] != '1')
          add(D[i][j][k + 1], D[i][j][k]);
        else
          add(D[i - 1][j][k], D[i][j][k]);
      }
  }

  long long ans = 0;
  for (int l = 1; l <= L; l++) {
    memset(pin, -1, sizeof(pin));
    for (int i = 0; i <= L; i++) {
      int s = 0x7fffff, e = -1;
      for (auto t : chk[l][i]) {
        s = min(s, t.first);
        e = max(e, t.second);
      }
      for (int j = s; j <= e; j++)
        pin[j + zc[i]] = i;
    }

    long long ret = 0;
    for (int j = 0; j <= L; j++)
      if (pin[j] != -1) {
        int p = pin[j];
        int z = j - zc[p];
        int o = l - j - oc[p];
        if (o >= 0) {
          add(ans, R[p][z][o]);
          add(ret, R[p][z][o]);
        }
      }
  }

  printf("%lld\\n", ans);

  return 0;
}