Problem p02636 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;

#define forn(i, n) for (int i = 0; i < (int)(n); i++)
#define si(c) ((int)(c).size())
#define forsn(i, s, n) for (int i = (int)(s); i < ((int)n); i++)
#define dforsn(i, s, n) for (int i = (int)(n)-1; i >= ((int)s); i--)
#define all(c) (c).begin(), (c).end()
#define D(a) cerr << #a << "=" << a << endl;
#define pb push_back
#define eb emplace_back
#define mp make_pair

typedef long long int ll;
typedef vector<int> vi;
typedef pair<int, int> pii;
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

const int MOD = 998244353;
const int N = 310;

int n;
string s;
bool can[2][N][N];
int best[N][N], suf[N][2], dp[2][N][N];

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

int memo[N][N][N];
int go(int from, int zeros, int ones) {
  if (zeros < 0 || ones < 0)
    return 0;
  if (zeros + ones == 0)
    return 1;

  int &res = memo[from][zeros][ones];
  if (~res)
    return res;

  res = 0;
  if (s[from] == '0')
    add(res, go(from + 1, zeros, ones));
  else
    add(res, go(from, zeros - 1, ones));

  if (s[from] == '1')
    add(res, go(from + 1, zeros, ones));
  else
    add(res, go(from, zeros, ones - 1));
  return res;
}

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

  cin >> s;
  n = si(s);
  dforsn(i, 0, n) {
    forn(j, 2) {
      suf[i][j] = s[i] == '0' + j;
      suf[i][j] += suf[i + 1][j];
    }
  }

  memset(best, -1, sizeof best);

  auto cur = can[0], prv = can[1];
  cur[0][0] = 1;
  for (int take = 1; take < n; take++) {
    swap(cur, prv);
    for (int a = 0; a <= n; a++) {
      for (int b = 0; b <= n - a; b++) {
        cur[a][b] = 0;
      }
    }
    for (int a = 0; a <= n; a++) {
      for (int b = 0; b <= n - a; b++) {
        if (prv[a][b]) {
          if (a >= 1 && a + b >= 2)
            cur[a - 1][b] = 1;
          if (b >= 1 && a + b >= 2)
            cur[a][b - 1] = 1;
          int rest = n - take + 1 - a - b;
          int pos = n - rest;
          if (a + b > 0)
            cur[a][b] = 1;
          if (s[pos] == '0' && b > 0)
            cur[a + 1][b - 1] = 1;
          if (s[pos] == '1' && a > 0)
            cur[a - 1][b + 1] = 1;
          if (pos + 1 < n) {
            if (s[pos] == '0' || s[pos + 1] == '0')
              cur[a + 1][b] = 1;
            if (s[pos] == '1' || s[pos + 1] == '1')
              cur[a][b + 1] = 1;
          }
        }
      }
    }

    for (int a = 0; a <= n; a++) {
      for (int b = 0; b <= n - a; b++)
        if (cur[a][b]) {
          int rest = n - take - a - b;
          int pos = n - rest;

          int aa = suf[pos][0] + a;
          int bb = suf[pos][1] + b;
          best[aa][bb] = max(best[aa][bb], pos);
        }
    }
  }
  s += '$';
  memset(memo, -1, sizeof memo);
  ll ans = 0;
  forn(a, n + 1) forn(b, n + 1) if (best[a][b] != -1) {
    int pos = best[a][b];
    int ca = a - suf[pos][0];
    int cb = b - suf[pos][1];
    ans += go(pos, ca, cb);
  }
  cout << (ans + 1) % MOD << endl;

  return 0;
}


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

#define forn(i, n) for (int i = 0; i < (int)(n); i++)
#define si(c) ((int)(c).size())
#define forsn(i, s, n) for (int i = (int)(s); i < ((int)n); i++)
#define dforsn(i, s, n) for (int i = (int)(n)-1; i >= ((int)s); i--)
#define all(c) (c).begin(), (c).end()
#define D(a) cerr << #a << "=" << a << endl;
#define pb push_back
#define eb emplace_back
#define mp make_pair

typedef long long int ll;
typedef vector<int> vi;
typedef pair<int, int> pii;
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

const int MOD = 998244353;
const int N = 310;

int n;
string s;
bool can[2][N][N];
int best[N][N], suf[N][2], dp[2][N][N];

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

int memo[N][N][N];
int go(int from, int zeros, int ones) {
  if (zeros < 0 || ones < 0)
    return 0;
  if (zeros + ones == 0)
    return 1;

  int &res = memo[from][zeros][ones];
  if (~res)
    return res;

  res = 0;
  if (s[from] == '0')
    add(res, go(from + 1, zeros, ones));
  else
    add(res, go(from, zeros - 1, ones));

  if (s[from] == '1')
    add(res, go(from + 1, zeros, ones));
  else
    add(res, go(from, zeros, ones - 1));
  return res;
}

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

  cin >> s;
  n = si(s);
  dforsn(i, 0, n) {
    forn(j, 2) {
      suf[i][j] = s[i] == '0' + j;
      suf[i][j] += suf[i + 1][j];
    }
  }

  memset(best, -1, sizeof best);

  auto cur = can[0], prv = can[1];
  cur[0][0] = 1;
  for (int take = 1; take < n; take++) {
    swap(cur, prv);
    for (int a = 0; a <= n; a++) {
      for (int b = 0; b <= n - a; b++) {
        cur[a][b] = 0;
      }
    }
    for (int a = 0; a <= n; a++) {
      for (int b = 0; b <= n - a; b++) {
        if (prv[a][b]) {
          if (a >= 1 && a + b >= 2)
            cur[a - 1][b] = 1;
          if (b >= 1 && a + b >= 2)
            cur[a][b - 1] = 1;
          int rest = n - take + 1 - a - b;
          int pos = n - rest;
          if (pos < n) {
            if (a + b > 0)
              cur[a][b] = 1;
            if (s[pos] == '0' && b > 0)
              cur[a + 1][b - 1] = 1;
            if (s[pos] == '1' && a > 0)
              cur[a - 1][b + 1] = 1;
            if (pos + 1 < n) {
              if (s[pos] == '0' || s[pos + 1] == '0')
                cur[a + 1][b] = 1;
              if (s[pos] == '1' || s[pos + 1] == '1')
                cur[a][b + 1] = 1;
            }
          }
        }
      }
    }

    for (int a = 0; a <= n; a++) {
      for (int b = 0; b <= n - a; b++)
        if (cur[a][b]) {
          int rest = n - take - a - b;
          int pos = n - rest;

          int aa = suf[pos][0] + a;
          int bb = suf[pos][1] + b;
          best[aa][bb] = max(best[aa][bb], pos);
        }
    }
  }
  s += '$';
  memset(memo, -1, sizeof memo);
  ll ans = 0;
  forn(a, n + 1) forn(b, n + 1) if (best[a][b] != -1) {
    int pos = best[a][b];
    int ca = a - suf[pos][0];
    int cb = b - suf[pos][1];
    ans += go(pos, ca, cb);
  }
  cout << (ans + 1) % MOD << endl;

  return 0;
}