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

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <string>
#include <vector>
template <typename T> void fin(T const &t) {
  std::cout << t << std::endl;
  exit(0);
}

template <typename T>
std::vector<std::vector<T>> mk2d(size_t s1, size_t s2, T const &val) {
  std::vector<T> tmp(s2, val);
  return std::vector<std::vector<T>>(s1, tmp);
}

template <typename T> std::vector<T> mk1d(size_t s, T const &val) {
  return std::vector<T>(s, val);
}

template <int64_t Prime> struct mod_int {
  int64_t v_;
  mod_int(int64_t x = 0) : v_(x) { normalize(); }
  void normalize() { ((v_ %= Prime) += Prime) %= Prime; }
  mod_int operator+=(mod_int const &r) {
    (v_ += r.v_) %= Prime;
    return *this;
  }
  mod_int operator-=(mod_int const &r) {
    (v_ += Prime - r.v_) %= Prime;
    return *this;
  }
  mod_int operator*=(mod_int const &r) {
    (v_ *= r.v_) %= Prime;
    return *this;
  }
  mod_int operator+(mod_int const &r) const {
    mod_int res(*this);
    return res += r;
  }
  mod_int operator-(mod_int const &r) const {
    mod_int res(*this);
    return res -= r;
  }
  mod_int operator*(mod_int const &r) const {
    mod_int res(*this);
    return res *= r;
  }
  mod_int pow(int x) const {
    int64_t res = 1, v = v_;
    while (x > 0) {
      if (x & 1)
        (res *= v) %= Prime;
      x /= 2;
      (v *= v) %= Prime;
    }
    return mod_int(res);
  }
  mod_int inv() const { return pow(Prime - 2); }
};

int64_t const MOD = 998244353;
typedef mod_int<MOD> mint;

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  std::string S;
  int K;
  std::cin >> S >> K;
  std::vector<int> a;
  int sum_one = 0, k = 0;
  for (auto it = S.rbegin(); it != S.rend(); ++it) {
    if (*it == '1')
      ++sum_one;
    else if (sum_one > 0) {
      a.push_back(sum_one);
      k += sum_one;
    }
  }
  if (K > sum_one * 2)
    K = sum_one * 2;
  auto tmp = mk2d(K + 1, sum_one + 1, mint(0));
  auto dp = mk1d(a.size() + 1, tmp);
  dp[0][0][0] = 1;
  for (int i = 0; i < a.size(); ++i) {
    std::vector<std::vector<mint>> &cur = dp[i + 1];
    std::vector<std::vector<mint>> &pre = dp[i];
    int ai1 = 0;
    if (i > 0)
      ai1 = a[i - 1];
    for (int l = 0; l <= ai1; ++l) {
      int p = ai1 - l;
      for (int k = 0; k <= K; ++k) {
        int ma = std::min(K - k + p, a[i] - l);
        for (int j = 0; j <= std::min(p, ma); ++j)
          cur[k][a[i] - j] += pre[k][l];
        for (int j = p + 1; j <= ma; ++j)
          cur[k + j - p][a[i] - j] += pre[k][l];
      }
    }
  }
  mint ans;
  for (int k = 0; k <= K; ++k)
    for (int l = 0; l <= sum_one; ++l)
      ans += dp[a.size()][k][l];
  fin(ans.v_);
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <string>
#include <vector>
template <typename T> void fin(T const &t) {
  std::cout << t << std::endl;
  exit(0);
}

template <typename T>
std::vector<std::vector<T>> mk2d(size_t s1, size_t s2, T const &val) {
  std::vector<T> tmp(s2, val);
  return std::vector<std::vector<T>>(s1, tmp);
}

template <typename T> std::vector<T> mk1d(size_t s, T const &val) {
  return std::vector<T>(s, val);
}

template <int64_t Prime> struct mod_int {
  int64_t v_;
  mod_int(int64_t x = 0) : v_(x) { normalize(); }
  void normalize() { ((v_ %= Prime) += Prime) %= Prime; }
  mod_int operator+=(mod_int const &r) {
    (v_ += r.v_) %= Prime;
    return *this;
  }
  mod_int operator-=(mod_int const &r) {
    (v_ += Prime - r.v_) %= Prime;
    return *this;
  }
  mod_int operator*=(mod_int const &r) {
    (v_ *= r.v_) %= Prime;
    return *this;
  }
  mod_int operator+(mod_int const &r) const {
    mod_int res(*this);
    return res += r;
  }
  mod_int operator-(mod_int const &r) const {
    mod_int res(*this);
    return res -= r;
  }
  mod_int operator*(mod_int const &r) const {
    mod_int res(*this);
    return res *= r;
  }
  mod_int pow(int x) const {
    int64_t res = 1, v = v_;
    while (x > 0) {
      if (x & 1)
        (res *= v) %= Prime;
      x /= 2;
      (v *= v) %= Prime;
    }
    return mod_int(res);
  }
  mod_int inv() const { return pow(Prime - 2); }
};

int64_t const MOD = 998244353;
typedef mod_int<MOD> mint;

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  std::string S;
  int K;
  std::cin >> S >> K;
  std::vector<int> a;
  int sum_one = 0, k = 0;
  for (auto it = S.rbegin(); it != S.rend(); ++it) {
    if (*it == '1')
      ++sum_one;
    else if (sum_one > 0) {
      a.push_back(sum_one);
      k += sum_one;
    }
  }
  if (K > k)
    K = k;
  auto tmp = mk2d(K + 1, sum_one + 1, mint(0));
  auto dp = mk1d(a.size() + 1, tmp);
  dp[0][0][0] = 1;
  for (int i = 0; i < a.size(); ++i) {
    std::vector<std::vector<mint>> &cur = dp[i + 1];
    std::vector<std::vector<mint>> &pre = dp[i];
    int ai1 = 0;
    if (i > 0)
      ai1 = a[i - 1];
    for (int l = 0; l <= ai1; ++l) {
      int p = ai1 - l;
      for (int k = 0; k <= K; ++k) {
        int ma = std::min(K - k + p, a[i] - l);
        for (int j = 0; j <= std::min(p, ma); ++j)
          cur[k][a[i] - j] += pre[k][l];
        for (int j = p + 1; j <= ma; ++j)
          cur[k + j - p][a[i] - j] += pre[k][l];
      }
    }
  }
  mint ans;
  for (int k = 0; k <= K; ++k)
    for (int l = 0; l <= sum_one; ++l)
      ans += dp[a.size()][k][l];
  fin(ans.v_);
  return 0;
}