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

#ifdef LOCAL
// #define _GLIBCXX_DEBUG
#endif

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

const int mod = 998244353;

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

void sub(int &a, int b) {
  a -= b;
  if (a < 0)
    a += mod;
}

int mult(int a, int b) { return ll(a) * b % mod; }

vector<vector<vector<int>>> calc(const string &s) {
  const int n = int(s.size());

  vector<vector<vector<int>>> ans(n + 1);
  for (int i = 0; i <= n; i++) {
    ans[i].resize(n - i + 1);
    for (int j = 0; j <= n - i; j++)
      ans[i][j].resize(n - i - j + 1);
  }

  const int N = 300;
  static int dp[N + 1][N + 1][N + 1][2];
  memset(dp, 0, sizeof dp);
  dp[0][0][0][0] = dp[0][0][0][1] = 1;

  for (int len = 0; len <= n; len++)
    for (int zero = 0; zero <= n - len; zero++)
      for (int one = 0; one <= n - len - zero; one++)
        for (int fin = 0; fin < 2; fin++) {
          const int cur = dp[len][zero][one][fin];

          if (fin == 1)
            add(ans[len][zero][one], cur);
          if (cur == 0 || len + zero + one >= n)
            continue;

          if (fin == 1) {
            add(dp[len][zero + 1][one][1], cur);
            add(dp[len][zero][one + 1][1], cur);
          }

          if (fin == 0) {
            add(dp[len + 1][zero][one][0], cur);
            add(dp[len + 1][zero][one][1], cur);

            if (s[len] == '0')
              add(dp[len][zero][one + 1][0], cur);
            else
              add(dp[len][zero + 1][one][0], cur);
          }
        }

  return ans;
}

void solve(istream &cin = std::cin, ostream &cout = std::cout) {
  string s;

  while (cin >> s) {
    const int n = int(s.size());

    string rev_s = s;
    reverse(rev_s.begin(), rev_s.end());

    auto dp_ways = calc(rev_s);

    cerr << double(clock()) / CLOCKS_PER_SEC << endl;

    vector<vector<vector<bool>>> reach(
        n + 1, vector<vector<bool>>(n + 1, vector<bool>(n + 1)));
    reach[0][0][0] = true;

    for (int i = 0; i <= n; i++) {
      for (int j = n; j >= 0; j--)
        for (int k = n; k >= 0; k--)
          if (reach[i][j][k]) {
            if (j >= 2)
              reach[i][j - 1][k] = true;
            if (k >= 2)
              reach[i][j][k - 1] = true;
            if (j >= 1 && k >= 1)
              reach[i][j][k - 1] = reach[i][j - 1][k] = true;
          }

      if (i + 1 <= n) {
        for (int j = 0; j <= n; j++)
          for (int k = 0; k <= n; k++)
            if (reach[i][j][k]) {
              const int nj = j + (s[i] == '0'), nk = k + (s[i] == '1');

              if (nj > 0)
                reach[i + 1][nj - 1][nk] = true;
              if (nk > 0)
                reach[i + 1][nj][nk - 1] = true;
            }
      }

      if (i + 2 <= n) {
        for (int j = 0; j <= n; j++)
          for (int k = 0; k <= n; k++)
            if (reach[i][j][k]) {
              for (int x = 0; x < 2; x++)
                if (s[i] == x + '0' || s[i + 1] == x + '0')
                  (x == 0 ? reach[i + 2][j + 1][k] : reach[i + 2][j][k + 1]) =
                      true;
            }
      }
    }

    int ans = 0;
    vector<int> cnt(2);
    vector<vector<bool>> used(n + 1, vector<bool>(n + 1, false));

    for (int suff = 0; suff <= n; suff++) {
      for (int zero = 0; zero <= n; zero++)
        for (int one = 0; one <= n; one++)
          if (reach[n - suff][zero][one]) {
            if (!used[cnt[0] + zero][cnt[1] + one]) {
              used[cnt[0] + zero][cnt[1] + one] = true;
              add(ans, dp_ways[suff][zero][one]);
            }
          }

      if (suff < n)
        cnt[rev_s[suff] - '0']++;
    }

    sub(ans, 1);
    cout << ans << endl;
  }
}

void gen(stringstream &cout) {
  const int n = 300;

  static mt19937 rng;
  string ans(n, '-');
  for (auto &ch : ans)
    ch = rng() % 2 + '0';

  cout << ans << endl;
}

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

#ifdef LOCAL
  auto st = clock();

  //	stringstream test, out;
  //	gen(test);
  //	solve(test, out);

  ifstream fin("../input.txt");
  solve(fin);

  cout << "clock: " << (clock() - st) / (double)CLOCKS_PER_SEC << endl;
#else
  solve();
#endif

  return 0;
}

Error: TLE

Fixed submission:

#ifdef LOCAL
// #define _GLIBCXX_DEBUG
#endif

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

const int mod = 998244353;

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

void sub(int &a, int b) {
  a -= b;
  if (a < 0)
    a += mod;
}

int mult(int a, int b) { return ll(a) * b % mod; }

vector<vector<vector<int>>> calc(const string &s) {
  const int n = int(s.size());

  vector<vector<vector<int>>> ans(
      n + 1, vector<vector<int>>(n + 1, vector<int>(n + 1)));

  vector<vector<vector<vector<int>>>> dp(
      n + 1, vector<vector<vector<int>>>(
                 n + 1, vector<vector<int>>(n + 1, vector<int>(2))));
  dp[0][0][0][0] = dp[0][0][0][1] = 1;

  for (int len = 0; len <= n; len++)
    for (int zero = 0; zero <= n; zero++)
      for (int one = 0; one <= n; one++)
        for (int fin = 0; fin < 2; fin++) {
          const int cur = dp[len][zero][one][fin];

          if (fin == 1)
            add(ans[len][zero][one], cur);
          if (cur == 0 || len + zero + one >= n)
            continue;

          for (int next_char = 0; next_char < 2; next_char++) {
            if (fin == 1)
              add(dp[len][zero + (next_char == 0)][one + (next_char == 1)][1],
                  cur);

            if (fin == 0 && next_char == s[len] - '0')
              for (int new_fin = fin; new_fin < 2; new_fin++)
                add(dp[len + 1][zero][one][new_fin], cur);

            if (fin == 0 && next_char != s[len] - '0')
              add(dp[len][zero + (next_char == 0)][one + (next_char == 1)][0],
                  cur);
          }
        }

  return ans;
}

void solve(istream &cin = std::cin, ostream &cout = std::cout) {
  string s;

  while (cin >> s) {
    const int n = int(s.size());

    string rev_s = s;
    reverse(rev_s.begin(), rev_s.end());

    auto dp_ways = calc(rev_s);

    //		for (int i = 0; i < int(dp_ways.size()); i++)
    //			for (int j = 0; j < int(dp_ways[i].size()); j++)
    //				for (int k = 0; k < int(dp_ways[i][j].size());
    //k++) 					if (dp_ways[i][j][k]) 						cerr << i << ' ' << j << ' ' << k << ' ' <<
    //dp_ways[i][j][k] << endl; 		break;

    vector<vector<vector<bool>>> reach(
        n + 1, vector<vector<bool>>(n + 1, vector<bool>(n + 1)));
    reach[0][0][0] = true;

    for (int i = 0; i <= n; i++) {
      for (int j = n; j >= 0; j--)
        for (int k = n; k >= 0; k--)
          if (reach[i][j][k]) {
            if (j >= 2)
              reach[i][j - 1][k] = true;
            if (k >= 2)
              reach[i][j][k - 1] = true;
            if (j >= 1 && k >= 1)
              reach[i][j][k - 1] = reach[i][j - 1][k] = true;
          }

      if (i + 1 <= n) {
        for (int j = 0; j <= n; j++)
          for (int k = 0; k <= n; k++)
            if (reach[i][j][k]) {
              const int nj = j + (s[i] == '0'), nk = k + (s[i] == '1');

              if (nj > 0)
                reach[i + 1][nj - 1][nk] = true;
              if (nk > 0)
                reach[i + 1][nj][nk - 1] = true;
            }
      }

      if (i + 2 <= n) {
        for (int j = 0; j <= n; j++)
          for (int k = 0; k <= n; k++)
            if (reach[i][j][k]) {
              for (int x = 0; x < 2; x++)
                if (s[i] == x + '0' || s[i + 1] == x + '0')
                  (x == 0 ? reach[i + 2][j + 1][k] : reach[i + 2][j][k + 1]) =
                      true;
            }
      }
    }

    int ans = 0;
    vector<int> cnt(2);
    vector<vector<bool>> used(n + 1, vector<bool>(n + 1, false));

    for (int suff = 0; suff <= n; suff++) {
      for (int zero = 0; zero <= n; zero++)
        for (int one = 0; one <= n; one++)
          if (reach[n - suff][zero][one]) {
            if (!used[cnt[0] + zero][cnt[1] + one]) {
              used[cnt[0] + zero][cnt[1] + one] = true;
              add(ans, dp_ways[suff][zero][one]);
            }
          }

      if (suff < n)
        cnt[rev_s[suff] - '0']++;
    }

    sub(ans, 1);
    cout << ans << endl;
  }
}

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

#ifdef LOCAL
  auto st = clock();

  ifstream fin("../input.txt");
  solve(fin);

  cout << "clock: " << (clock() - st) / (double)CLOCKS_PER_SEC << endl;
#else
  solve();
#endif

  return 0;
}