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

#include <bits/stdc++.h>
#define ld long double
#define ll long long
#define st first
#define nd second
#define mp make_pair
#define pii pair<int, int>
#define vi vector<int>
#define pb push_back
using namespace std;
const ll mod = 998244353;
const int MAXN = 301;
int n, k;
string s;
ll dp[MAXN][MAXN], f[MAXN];
ll inv(ll x, ll p = mod - 2) {
  if (p == 0)
    return 1;
  ll y = inv(x, p / 2);
  y = y * y % mod;
  if (p % 2 == 0) {
    return y;
  }
  return y * x % mod;
}
ll C(int N, int K) {
  if (K > N)
    return 1;
  return f[N] * inv(f[K] * f[N - K] % mod) % mod;
}
int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  cin >> s >> k;
  n = s.size();
  f[0] = 1;
  for (int i = 1; i <= n; ++i) {
    f[i] = f[i - 1] * (ll)i % mod;
  }
  k = min(k, n);
  vi block;
  if (s[0] == '0')
    block.push_back(1);
  else {
    block.pb(0);
    block.pb(1);
  }
  for (int i = 1; i < n; ++i) {
    if (s[i] != s[i - 1]) {
      block.pb(0);
    }
    block.back()++;
  }
  dp[0][0] = 1;
  for (int i = (int)block.size() - 1; i >= 0; --i) {
    int cnt = block[i];
    if (cnt == 0)
      continue;
    vector<vector<ll>> prefj(k + 1, vector<ll>(k + 1));
    for (int x = 0; x <= k; ++x) {
      prefj[x][0] = dp[x][0];
      for (int j = 1; j <= k; ++j) {
        prefj[x][j] = (prefj[x][j - 1] + dp[x][j]) % mod;
      }
    }
    if (i % 2 == 0) {
      for (int c = 1; c < cnt; ++c) {
        for (int j = k; j >= 0; --j) {
          vector<ll> pref(k + 1);
          pref[0] = dp[0][j];
          for (int x = 1; x <= j; ++x) {
            pref[x] = (pref[x - 1] + dp[x][j]) % mod;
          }
          for (int x = j; x >= 0; --x) {
            dp[x][j] += ((x == 0) ? 0 : pref[x - 1]);
            dp[x][j] %= mod;
          }
        }
      }
    }

    if (i % 2 == 1) {
      for (int j = k; j >= 0; --j) {
        vector<ll> pref(k + 1);
        pref[0] = dp[0][j];
        for (int x = 1; x <= j; ++x) {
          pref[x] = (pref[x - 1] + dp[x][j]) % mod;
        }
        for (int x = j; x >= 0; --x) {
          if ((int)block.size() - i > 2) {
            dp[x][j] += ((x == 0) ? 0 : pref[x - 1]);
            dp[x][j] %= mod;
          }
          if (j != 0) {
            dp[x][j] +=
                (prefj[x][j - 1] -
                 ((max(0, j - cnt) == 0) ? 0 : prefj[x][j - cnt - 1]) + mod) %
                mod;
            dp[x][j] %= mod;
          }
        }
      }
    }
  }
  ll ans = 0;
  for (int i = 0; i <= k; ++i) {
    ans += dp[i][i];
    ans %= mod;
    for (int j = 0; j < i; ++j) {
      if (block[0] > 0) {
        ans += dp[j][i];
        ans %= mod;
      }
    }
  }
  cout << ans;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define ld long double
#define ll long long
#define st first
#define nd second
#define mp make_pair
#define pii pair<int, int>
#define vi vector<int>
#define pb push_back
using namespace std;
const ll mod = 998244353;
const int MAXN = 301;
int n, k;
string s;
ll dp[MAXN][MAXN], f[MAXN];
ll inv(ll x, ll p = mod - 2) {
  if (p == 0)
    return 1;
  ll y = inv(x, p / 2);
  y = y * y % mod;
  if (p % 2 == 0) {
    return y;
  }
  return y * x % mod;
}
ll C(int N, int K) {
  if (K > N)
    return 1;
  return f[N] * inv(f[K] * f[N - K] % mod) % mod;
}
int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  cin >> s >> k;
  n = s.size();
  f[0] = 1;
  for (int i = 1; i <= n; ++i) {
    f[i] = f[i - 1] * (ll)i % mod;
  }
  k = min(k, n);
  vi block;
  if (s[0] == '0')
    block.push_back(1);
  else {
    block.pb(0);
    block.pb(1);
  }
  for (int i = 1; i < n; ++i) {
    if (s[i] != s[i - 1]) {
      block.pb(0);
    }
    block.back()++;
  }
  dp[0][0] = 1;
  for (int i = (int)block.size() - 1; i >= 0; --i) {
    int cnt = block[i];
    if (cnt == 0)
      continue;
    vector<vector<ll>> prefj(k + 1, vector<ll>(k + 1));
    for (int x = 0; x <= k; ++x) {
      prefj[x][0] = dp[x][0];
      for (int j = 1; j <= k; ++j) {
        prefj[x][j] = (prefj[x][j - 1] + dp[x][j]) % mod;
      }
    }
    for (int j = k; j >= 0; --j) {
      vector<ll> pref(j + 1), prefC(j + 1);
      pref[0] = dp[0][j];
      prefC[0] = dp[0][j] * C(0, cnt - 1) % mod;
      for (int x = 1; x <= j; ++x) {
        pref[x] = (pref[x - 1] + dp[x][j]) % mod;
        prefC[x] = (prefC[x - 1] + dp[x][j] * C(x + cnt - 2, x)) % mod;
      }
      for (int x = j; x >= 0; --x) {
        if (i % 2 == 0) { // '0'
          if (cnt > 1) {
            //	dp[x][j] +=	((x==0)?0:prefC[x-1]);
            //	dp[x][j] %= mod;
            for (int x1 = 0; x1 < x; ++x1) {
              dp[x][j] += dp[x1][j] * C((x - x1) + cnt - 2, (x - x1)) % mod;
              dp[x][j] %= mod;
            }
          }
        } else { // '1'
          if ((int)block.size() - i > 2) {
            dp[x][j] += ((x == 0) ? 0 : pref[x - 1]);
            dp[x][j] %= mod;
          }
          if (j != 0) {
            dp[x][j] +=
                (prefj[x][j - 1] -
                 ((max(0, j - cnt) == 0) ? 0 : prefj[x][j - cnt - 1]) + mod) %
                mod;
            dp[x][j] %= mod;
          }
        }
      }
    }
  }
  ll ans = 0;
  for (int i = 0; i <= k; ++i) {
    ans += dp[i][i];
    ans %= mod;
    for (int j = 0; j < i; ++j) {
      if (block[0] > 0) {
        ans += dp[j][i];
        ans %= mod;
      }
    }
  }
  cout << ans;
}