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

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define chmax(a, b) a = max(a, b);
#define chmin(a, b) a = min(a, b);
using namespace std;
using ll = long long;
using P = pair<int, int>;

// #define MOD 1000000007
#define MOD 998244353

class mint {
public:
  int i;

  mint() : i(0) {}

  template <class T> mint(T x) {
    i = int(x % MOD);
    if (i < 0)
      i += MOD;
  }

  mint operator+(mint x) { return i + x.i; }

  mint operator-(mint x) { return i - x.i; }

  mint operator*(mint x) { return (long long)i * x.i; }

  mint operator/(mint x) { return (long long)i * x.pow(MOD - 2).i; }

  mint inv() { return pow(MOD - 2); }

  template <class T> mint pow(T p) {
    long long r = 1;
    long long t = i;
    for (; p > 0; p >>= 1) {
      if (p & 1)
        r = r * t % MOD;
      t = t * t % MOD;
    }
    return r;
  }

  template <class T1, class T2> static mint pow(T1 a, T2 b) {
    long long r = 1;
    long long t = (long long)(a % MOD);
    for (; b > 0; b >>= 1) {
      if (b & 1)
        r = r * t % MOD;
      t = t * t % MOD;
    }
    return r;
  }

  mint &operator+=(mint x) {
    i = (i + x.i) % MOD;
    return *this;
  }

  mint &operator-=(mint x) {
    i = i - x.i;
    if (i < 0)
      i += MOD;
    return *this;
  }

  mint &operator*=(mint x) {
    i = (int)((long long)i * x.i % MOD);
    return *this;
  }

  mint &operator/=(mint x) {
    i = (long long)i * x.pow(MOD - 2).i % MOD;
    return *this;
  }
};

std::ostream &operator<<(std::ostream &os, const mint &m) { return os << m.i; }

int main() {
  string s;
  int k;
  cin >> s >> k;

  vector<int> d;
  {
    int pre = 0;
    rep(i, s.size()) {
      if (s[i] == '0') {
        d.push_back(i - pre);
        pre = i + 1;
      }
    }
    if (pre != s.size())
      d.push_back(s.size() - pre);
  }
  reverse(d.begin(), d.end());
  k = min(k, 305);
  vector<vector<vector<mint>>> dp(
      d.size() + 1, vector<vector<mint>>(k + 1, vector<mint>(k + 1)));
  dp[0][0][k] = 1;

  rep(i, d.size()) rep(sc, k + 1) rep(lft, k + 1) {
    int cnt = d[i];
    mint v = dp[i][sc][lft];
    if (v.i == 0)
      continue;
    for (int r = 1; r <= cnt; r++) {
      if (!(lft - r >= 0 && sc + r <= k))
        break;
      dp[i + 1][sc + r][lft - r] += v;
    }
    dp[i + 1][sc][lft] += v;
    for (int r = 1; r <= sc; r++) {
      dp[i + 1][sc - r][lft] += v;
    }
  }
  mint ans = 0;
  rep(i, k + 1) ans += dp[d.size()][0][i];
  cout << ans << endl;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define chmax(a, b) a = max(a, b);
#define chmin(a, b) a = min(a, b);
using namespace std;
using ll = long long;
using P = pair<int, int>;

// #define MOD 1000000007
#define MOD 998244353

class mint {
public:
  int i;

  mint() : i(0) {}

  template <class T> mint(T x) {
    i = int(x % MOD);
    if (i < 0)
      i += MOD;
  }

  mint operator+(mint x) { return i + x.i; }

  mint operator-(mint x) { return i - x.i; }

  mint operator*(mint x) { return (long long)i * x.i; }

  mint operator/(mint x) { return (long long)i * x.pow(MOD - 2).i; }

  mint inv() { return pow(MOD - 2); }

  template <class T> mint pow(T p) {
    long long r = 1;
    long long t = i;
    for (; p > 0; p >>= 1) {
      if (p & 1)
        r = r * t % MOD;
      t = t * t % MOD;
    }
    return r;
  }

  template <class T1, class T2> static mint pow(T1 a, T2 b) {
    long long r = 1;
    long long t = (long long)(a % MOD);
    for (; b > 0; b >>= 1) {
      if (b & 1)
        r = r * t % MOD;
      t = t * t % MOD;
    }
    return r;
  }

  mint &operator+=(mint x) {
    i = (i + x.i) % MOD;
    return *this;
  }

  mint &operator-=(mint x) {
    i = i - x.i;
    if (i < 0)
      i += MOD;
    return *this;
  }

  mint &operator*=(mint x) {
    i = (int)((long long)i * x.i % MOD);
    return *this;
  }

  mint &operator/=(mint x) {
    i = (long long)i * x.pow(MOD - 2).i % MOD;
    return *this;
  }
};

std::ostream &operator<<(std::ostream &os, const mint &m) { return os << m.i; }

int main() {
  string s;
  int k;
  cin >> s >> k;

  vector<int> d;
  {
    int pre = 0;
    rep(i, s.size()) {
      if (s[i] == '0') {
        d.push_back(i - pre);
        pre = i + 1;
      }
    }
    if (pre != s.size())
      d.push_back(s.size() - pre);
  }
  reverse(d.begin(), d.end());
  k = min(k, 305);
  vector<vector<vector<mint>>> dp(
      d.size() + 1, vector<vector<mint>>(k + 1, vector<mint>(k + 1)));
  dp[0][0][k] = 1;

  rep(i, d.size()) rep(sc, k + 1) rep(lft, k + 1) {
    int cnt = d[i];
    for (int r = 1; r <= cnt; r++) {
      if (!(lft - r >= 0 && sc + r <= k))
        break;
      dp[i + 1][sc + r][lft - r] += dp[i][sc][lft];
    }
    dp[i + 1][sc][lft] += dp[i][sc][lft];
    for (int r = 1; r <= sc; r++) {
      dp[i + 1][sc - r][lft] += dp[i][sc][lft];
    }
  }
  mint ans = 0;
  rep(i, k + 1) ans += dp[d.size()][0][i];
  cout << ans << endl;
}