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

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

#define rep(i, a, b) for (int i = a; i < (b); ++i)
#define trav(a, x) for (auto &a : (x))
#define all(x) x.begin(), x.end() // sort(all(x))
#define sz(x) (int)(x).size()
#define mp(x, y) make_pair((x), (y))

typedef long long ll;

typedef double T;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef tuple<int, int, int> tiii;

typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vi> vvi;
typedef vector<vl> vvl;
typedef vector<T> vt;
typedef vector<vt> vvt;
typedef vector<pii> vpi;
typedef vector<pii> vpl;
typedef vector<vpi> graph;
typedef complex<T> Pt;
typedef pair<Pt, Pt> Line;

const ll MOD = 998244353;
const T EPS = 1e-10;
const int INF = 1e9 + 1;

/* Mod Library */

struct Mod {
  ll x;
  Mod(ll x = 0) : x(x) {}
  Mod operator+(Mod b) { return Mod((x + b.x) % MOD); }
  const Mod &operator+=(Mod b) {
    x = (x + b.x) % MOD;
    return *this;
  }
  Mod operator-(Mod b) { return Mod((x - b.x + MOD) % MOD); }
  const Mod &operator-=(Mod b) {
    x = (x - b.x + MOD) % MOD;
    return *this;
  }
  Mod operator*(Mod b) { return Mod((x * b.x) % MOD); }
  Mod operator/(Mod b) { return *this * inverse(b); }
  Mod operator%(Mod b) { return Mod(x % b.x); }

  static ll euclid(ll a, ll b, ll &x, ll &y) {
    if (b) {
      ll d = euclid(b, a % b, y, x);
      y -= a / b * x;
      return d;
    }
    x = 1;
    y = 0;
    return a;
  }

  Mod inverse(Mod a) {
    ll x, y, g = euclid(a.x, MOD, x, y);
    assert(g == 1);
    return Mod((x + MOD) % MOD);
  }

  Mod operator^(ll e) {
    if (!e)
      return Mod(1);
    Mod r = *this ^ (e / 2);
    r = r * r;
    return e & 1 ? *this * r : r;
  }
};

istream &operator>>(istream &in, Mod &b) {
  in >> b.x;
  return in;
}

ostream &operator<<(ostream &out, Mod &b) {
  out << b.x;
  return out;
}

/* Input Library */
template <typename A, typename B>
istream &operator>>(istream &in, pair<A, B> &p) {
  in >> p.first >> p.second;
  return in;
}

template <typename A, typename B>
ostream &operator<<(ostream &out, pair<A, B> &p) {
  out << p.first << " " << p.second;
  return out;
}

template <typename I> istream &operator>>(istream &in, vector<I> &a) {
  trav(i, a) in >> i;
  return in;
}

template <typename I> ostream &operator<<(ostream &out, vector<I> &a) {
  rep(i, 0, sz(a)) {
    if (i)
      out << " ";
    out << a[i];
  }
  return out;
}

void fast_io() {
  cin.sync_with_stdio(0);
  cin.tie(0);
  // cout.tie(0);
  cin.exceptions(cin.failbit);
}

// index * num moves used * num ones available.
Mod dp[301][301][301];

void solve_test() {
  string ss;
  int n, K;
  cin >> ss >> K;
  ss = ss + '0';
  while (ss.size() && ss[0] == '1') {
    ss = ss.substr(1);
  }
  n = ss.size();
  K = min(K, n);
  if (K == 0) {
    cout << 1 << endl;
    return;
  }

  vector<int> s;
  for (int i = 0; i < n; ++i) {
    int tl = i;
    while (tl < n && ss[tl] == '1') {
      ++tl;
    }
    s.emplace_back(tl - i);
    i = tl;
  }
  n = s.size();

  dp[n][0][0] = 1;
  for (int i = n - 1; i >= 0; --i) {
    for (int j = 0; j <= K; ++j) {
      for (int k = 0; k <= K - j; ++k) {
        // giving up non-zero amount of ones.
        for (int l = max(s[i] - k, 0); l <= s[i]; ++l) {
          dp[i][j][k] += dp[i + 1][j][k - s[i] + l];
        }
        // taking positive amount of ones.
        for (int l = 1; l <= min(j, K - k); ++l) {
          dp[i][j][k] += dp[i + 1][j - l][k + l];
        }
      }
    }
  }

  Mod ans = 0;
  for (int i = 0; i <= K; ++i) {
    ans += dp[0][i][0];
  }
  cout << ans << endl;
}

void setup() {}

int main() {
  setup();
  srand(time(0));
  fast_io();
  int t = 1;
  // cin >> t;
  while (t--) {
    solve_test();
  }
}


Error: TLE

Fixed submission:

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

#define rep(i, a, b) for (int i = a; i < (b); ++i)
#define trav(a, x) for (auto &a : (x))
#define all(x) x.begin(), x.end() // sort(all(x))
#define sz(x) (int)(x).size()
#define mp(x, y) make_pair((x), (y))

typedef long long ll;

typedef double T;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef tuple<int, int, int> tiii;

typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vi> vvi;
typedef vector<vl> vvl;
typedef vector<T> vt;
typedef vector<vt> vvt;
typedef vector<pii> vpi;
typedef vector<pii> vpl;
typedef vector<vpi> graph;
typedef complex<T> Pt;
typedef pair<Pt, Pt> Line;

const ll MOD = 998244353;
const T EPS = 1e-10;
const int INF = 1e9 + 1;

/* Mod Library */

struct Mod {
  ll x;
  Mod(ll x = 0) : x(x) {}
  Mod operator+(Mod b) { return Mod((x + b.x) % MOD); }
  const Mod &operator+=(Mod b) {
    x = (x + b.x) % MOD;
    return *this;
  }
  Mod operator-(Mod b) { return Mod((x - b.x + MOD) % MOD); }
  const Mod &operator-=(Mod b) {
    x = (x - b.x + MOD) % MOD;
    return *this;
  }
  Mod operator*(Mod b) { return Mod((x * b.x) % MOD); }
  Mod operator/(Mod b) { return *this * inverse(b); }
  Mod operator%(Mod b) { return Mod(x % b.x); }

  static ll euclid(ll a, ll b, ll &x, ll &y) {
    if (b) {
      ll d = euclid(b, a % b, y, x);
      y -= a / b * x;
      return d;
    }
    x = 1;
    y = 0;
    return a;
  }

  Mod inverse(Mod a) {
    ll x, y, g = euclid(a.x, MOD, x, y);
    assert(g == 1);
    return Mod((x + MOD) % MOD);
  }

  Mod operator^(ll e) {
    if (!e)
      return Mod(1);
    Mod r = *this ^ (e / 2);
    r = r * r;
    return e & 1 ? *this * r : r;
  }
};

istream &operator>>(istream &in, Mod &b) {
  in >> b.x;
  return in;
}

ostream &operator<<(ostream &out, Mod &b) {
  out << b.x;
  return out;
}

/* Input Library */
template <typename A, typename B>
istream &operator>>(istream &in, pair<A, B> &p) {
  in >> p.first >> p.second;
  return in;
}

template <typename A, typename B>
ostream &operator<<(ostream &out, pair<A, B> &p) {
  out << p.first << " " << p.second;
  return out;
}

template <typename I> istream &operator>>(istream &in, vector<I> &a) {
  trav(i, a) in >> i;
  return in;
}

template <typename I> ostream &operator<<(ostream &out, vector<I> &a) {
  rep(i, 0, sz(a)) {
    if (i)
      out << " ";
    out << a[i];
  }
  return out;
}

void fast_io() {
  cin.sync_with_stdio(0);
  cin.tie(0);
  // cout.tie(0);
  cin.exceptions(cin.failbit);
}

// index * num moves used * num ones available.
Mod dp[301][301][301];

void solve_test() {
  string ss;
  int n, K;
  cin >> ss >> K;
  ss = ss + '0';
  while (ss.size() && ss[0] == '1') {
    ss = ss.substr(1);
  }
  n = ss.size();
  K = min(K, n);
  if (K == 0) {
    cout << 1 << endl;
    return;
  }

  vector<int> s;
  for (int i = 0; i < n; ++i) {
    int tl = i;
    while (tl < n && ss[tl] == '1') {
      ++tl;
    }
    s.emplace_back(tl - i);
    i = tl;
  }
  n = s.size();

  dp[n][0][0] = 1;
  for (int i = n - 1; i >= 0; --i) {
    for (int j = 0; j <= K; ++j) {
      for (int k = 0; k <= K; ++k) {
        // giving up non-zero amount of ones.
        for (int l = max(s[i] - k, 0); l <= s[i]; ++l) {
          dp[i][j][k] += dp[i + 1][j][k - s[i] + l];
        }
        // taking positive amount of ones.
        for (int l = 1; l <= min(j, K - k); ++l) {
          dp[i][j][k] += dp[i + 1][j - l][k + l];
        }
      }
    }
  }

  Mod ans = 0;
  for (int i = 0; i <= K; ++i) {
    ans += dp[0][i][0];
  }
  cout << ans << endl;
}

void setup() {}

int main() {
  setup();
  srand(time(0));
  fast_io();
  int t = 1;
  // cin >> t;
  while (t--) {
    solve_test();
  }
}