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

#include <bits/stdc++.h>
using namespace std;
using LL = long long int;
#define incID(i, l, r) for (int i = (l); i < (r); ++i)
#define decID(i, l, r) for (int i = (r)-1; i >= (l); --i)
#define incII(i, l, r) for (int i = (l); i <= (r); ++i)
#define decII(i, l, r) for (int i = (r); i >= (l); --i)
#define inc(i, n) incID(i, 0, n)
#define dec(i, n) decID(i, 0, n)
#define inc1(i, n) incII(i, 1, n)
#define dec1(i, n) decII(i, 1, n)
#define inID(v, l, r) ((l) <= (v) && (v) < (r))
#define inII(v, l, r) ((l) <= (v) && (v) <= (r))
#define PB push_back
#define EB emplace_back
#define MP make_pair
#define MT make_tuple
#define FI first
#define SE second
#define FR front()
#define BA back()
#define ALL(v) v.begin(), v.end()
#define RALL(v) v.rbegin(), v.rend()
auto setmin = [](auto &a, auto b) { return (b < a ? a = b, true : false); };
auto setmax = [](auto &a, auto b) { return (b > a ? a = b, true : false); };
auto setmineq = [](auto &a, auto b) { return (b <= a ? a = b, true : false); };
auto setmaxeq = [](auto &a, auto b) { return (b >= a ? a = b, true : false); };
#define SI(v) static_cast<int>(v.size())
#define RF(e, v) for (auto &e : v)
#define until(e) while (!(e))
#define if_not(e) if (!(e))
#define ef else if
#define UR assert(false)
#define IN(T, ...)                                                             \\
  T __VA_ARGS__;                                                               \\
  IN_(__VA_ARGS__);
void IN_(){};
template <typename T, typename... U> void IN_(T &a, U &...b) {
  cin >> a;
  IN_(b...);
};
template <typename T> void OUT(T &&a) { cout << a << endl; }
template <typename T, typename... U> void OUT(T &&a, U &&...b) {
  cout << a << " ";
  OUT(b...);
}

// ---- ----

template <LL M> class ModInt {
private:
  LL v;
  pair<LL, LL> ext_gcd(LL a, LL b) {
    if (b == 0) {
      assert(a == 1);
      return {1, 0};
    }
    auto p = ext_gcd(b, a % b);
    return {p.SE, p.FI - (a / b) * p.SE};
  }

public:
  ModInt(LL vv = 0) {
    v = vv;
    if (abs(v) >= M) {
      v %= M;
    }
    if (v < 0) {
      v += M;
    }
  }
  LL get_v() { return v; }
  ModInt inv() { return ext_gcd(M, v).SE; }
  ModInt exp(LL b) {
    ModInt p = 1, a = v;
    if (b < 0) {
      a = a.inv();
      b = -b;
    }
    while (b) {
      if (b & 1) {
        p *= a;
      }
      a *= a;
      b >>= 1;
    }
    return p;
  }
  friend bool operator<(ModInt a, ModInt b) { return (a.v < b.v); }
  friend bool operator>(ModInt a, ModInt b) { return (a.v > b.v); }
  friend bool operator<=(ModInt a, ModInt b) { return (a.v <= b.v); }
  friend bool operator>=(ModInt a, ModInt b) { return (a.v >= b.v); }
  friend bool operator==(ModInt a, ModInt b) { return (a.v == b.v); }
  friend bool operator!=(ModInt a, ModInt b) { return (a.v != b.v); }
  friend ModInt operator+(ModInt a) { return ModInt(+a.v); }
  friend ModInt operator-(ModInt a) { return ModInt(-a.v); }
  friend ModInt operator+(ModInt a, ModInt b) { return ModInt(a.v + b.v); }
  friend ModInt operator-(ModInt a, ModInt b) { return ModInt(a.v - b.v); }
  friend ModInt operator*(ModInt a, ModInt b) { return ModInt(a.v * b.v); }
  friend ModInt operator/(ModInt a, ModInt b) { return a * b.inv(); }
  friend ModInt operator^(ModInt a, LL b) { return a.exp(b); }
  friend ModInt &operator+=(ModInt &a, ModInt b) { return (a = a + b); }
  friend ModInt &operator-=(ModInt &a, ModInt b) { return (a = a - b); }
  friend ModInt &operator*=(ModInt &a, ModInt b) { return (a = a * b); }
  friend ModInt &operator/=(ModInt &a, ModInt b) { return (a = a / b); }
  friend ModInt &operator^=(ModInt &a, LL b) { return (a = a ^ b); }
  friend istream &operator>>(istream &s, ModInt &b) {
    s >> b.v;
    b = ModInt(b.v);
    return s;
  }
  friend ostream &operator<<(ostream &s, ModInt b) { return (s << b.v); }
};

// ----

using MI = ModInt<998244353>;

template <typename T> T MV(T v) { return v; }
template <typename T, typename... U> auto MV(T v, int a, U... b) {
  return vector<decltype(MV(v, b...))>(a, MV(v, b...));
}

#define inCD(v, l, r) ((l) < (v) && (v) < (r))
#define inCI(v, l, r) ((l) < (v) && (v) <= (r))

int main() {
  IN(string, str);
  IN(int, K);
  vector<int> v = {0};
  int s = 0;
  RF(c, str) {
    if (c == '1') {
      v.BA++;
      s++;
    } else {
      v.PB(0);
    }
  }
  int n = SI(v);

  MI zero = 0;
  auto dp_ = MV<MI>(0, n + 1, s + 1, 2 * s + 1);
  auto dp = [&](int i, int j, int k) -> auto & {
    zero = 0;
    return (inII(j, 0, s) && inII(k, 0, 2 * s) ? dp_[i][j][k] : zero);
  };
  dp(0, 0, 0) += 1;
  dp(0, 1, 1) -= 1;
  incII(i, 0, n) {
    incII(k, 0, 2 * s) {
      incII(j, 0, s) {
        auto &x = dp(i, j, k);
        dp(i, j - 1, k + 1) += x;
        dp(i, j + 1, k + 1) += x;
        if (inCD(j, 0, s)) {
          dp(i, j, k + 2) -= x;
        }
        if (i == n) {
          continue;
        }

        dp(i + 1, j, k) += x;
        dp(i + 1, j, k + 2) -= x;
        if (j > v[i]) {
          dp(i + 1, j - v[i] - 1, k + v[i] + 1) -= x;
          dp(i + 1, j - v[i], k + v[i] + 2) += x;
        }
      }
    }
  }
  MI ans = 0;
  incII(k, 0, 2 * s) {
    if (k <= 2 * K) {
      ans += dp(n, 0, k);
    }
  }
  OUT(ans);
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using LL = long long int;
#define incID(i, l, r) for (int i = (l); i < (r); ++i)
#define decID(i, l, r) for (int i = (r)-1; i >= (l); --i)
#define incII(i, l, r) for (int i = (l); i <= (r); ++i)
#define decII(i, l, r) for (int i = (r); i >= (l); --i)
#define inc(i, n) incID(i, 0, n)
#define dec(i, n) decID(i, 0, n)
#define inc1(i, n) incII(i, 1, n)
#define dec1(i, n) decII(i, 1, n)
#define inID(v, l, r) ((l) <= (v) && (v) < (r))
#define inII(v, l, r) ((l) <= (v) && (v) <= (r))
#define PB push_back
#define EB emplace_back
#define MP make_pair
#define MT make_tuple
#define FI first
#define SE second
#define FR front()
#define BA back()
#define ALL(v) v.begin(), v.end()
#define RALL(v) v.rbegin(), v.rend()
auto setmin = [](auto &a, auto b) { return (b < a ? a = b, true : false); };
auto setmax = [](auto &a, auto b) { return (b > a ? a = b, true : false); };
auto setmineq = [](auto &a, auto b) { return (b <= a ? a = b, true : false); };
auto setmaxeq = [](auto &a, auto b) { return (b >= a ? a = b, true : false); };
#define SI(v) static_cast<int>(v.size())
#define RF(e, v) for (auto &e : v)
#define until(e) while (!(e))
#define if_not(e) if (!(e))
#define ef else if
#define UR assert(false)
#define IN(T, ...)                                                             \\
  T __VA_ARGS__;                                                               \\
  IN_(__VA_ARGS__);
void IN_(){};
template <typename T, typename... U> void IN_(T &a, U &...b) {
  cin >> a;
  IN_(b...);
};
template <typename T> void OUT(T &&a) { cout << a << endl; }
template <typename T, typename... U> void OUT(T &&a, U &&...b) {
  cout << a << " ";
  OUT(b...);
}

// ---- ----

template <LL M> class ModInt {
private:
  LL v;
  pair<LL, LL> ext_gcd(LL a, LL b) {
    if (b == 0) {
      assert(a == 1);
      return {1, 0};
    }
    auto p = ext_gcd(b, a % b);
    return {p.SE, p.FI - (a / b) * p.SE};
  }

public:
  ModInt(LL vv = 0) {
    v = vv;
    if (abs(v) >= M) {
      v %= M;
    }
    if (v < 0) {
      v += M;
    }
  }
  LL get_v() { return v; }
  ModInt inv() { return ext_gcd(M, v).SE; }
  ModInt exp(LL b) {
    ModInt p = 1, a = v;
    if (b < 0) {
      a = a.inv();
      b = -b;
    }
    while (b) {
      if (b & 1) {
        p *= a;
      }
      a *= a;
      b >>= 1;
    }
    return p;
  }
  friend bool operator<(ModInt a, ModInt b) { return (a.v < b.v); }
  friend bool operator>(ModInt a, ModInt b) { return (a.v > b.v); }
  friend bool operator<=(ModInt a, ModInt b) { return (a.v <= b.v); }
  friend bool operator>=(ModInt a, ModInt b) { return (a.v >= b.v); }
  friend bool operator==(ModInt a, ModInt b) { return (a.v == b.v); }
  friend bool operator!=(ModInt a, ModInt b) { return (a.v != b.v); }
  friend ModInt operator+(ModInt a) { return ModInt(+a.v); }
  friend ModInt operator-(ModInt a) { return ModInt(-a.v); }
  friend ModInt operator+(ModInt a, ModInt b) { return ModInt(a.v + b.v); }
  friend ModInt operator-(ModInt a, ModInt b) { return ModInt(a.v - b.v); }
  friend ModInt operator*(ModInt a, ModInt b) { return ModInt(a.v * b.v); }
  friend ModInt operator/(ModInt a, ModInt b) { return a * b.inv(); }
  friend ModInt operator^(ModInt a, LL b) { return a.exp(b); }
  friend ModInt &operator+=(ModInt &a, ModInt b) { return (a = a + b); }
  friend ModInt &operator-=(ModInt &a, ModInt b) { return (a = a - b); }
  friend ModInt &operator*=(ModInt &a, ModInt b) { return (a = a * b); }
  friend ModInt &operator/=(ModInt &a, ModInt b) { return (a = a / b); }
  friend ModInt &operator^=(ModInt &a, LL b) { return (a = a ^ b); }
  friend istream &operator>>(istream &s, ModInt &b) {
    s >> b.v;
    b = ModInt(b.v);
    return s;
  }
  friend ostream &operator<<(ostream &s, ModInt b) { return (s << b.v); }
};

// ----

using MI = ModInt<998244353>;

template <typename T> T MV(T v) { return v; }
template <typename T, typename... U> auto MV(T v, int a, U... b) {
  return vector<decltype(MV(v, b...))>(a, MV(v, b...));
}

int main() {
  IN(string, s);
  IN(int, k);

  s.PB('0');
  vector<int> v = {0};
  RF(c, s) {
    if (c == '1') {
      v.BA++;
    } else {
      v.PB(0);
    }
  }
  int n = SI(v), sum = 0;
  RF(e, v) { sum += e; }

  vector<map<pair<int, int>, MI>> dp(n + 1);
  dp[0][{0, 0}] += 1;
  dp[0][{-1, +1}] -= 1;
  dp[0][{+1, +1}] -= 1;
  dp[0][{0, +2}] += 1;
  inc(i, n + 1) {
    incII(l, 0, 2 * sum) {
      incII(j, -1, sum + 1) {
        dp[i][{j, l}] += dp[i][{j - 1, l - 1}] + dp[i][{j + 1, l - 1}];
        if_not(inII(j, 0, sum)) { continue; }
        dp[i][{j, l}] -= dp[i][{j, l - 2}];
        if (i == n) {
          continue;
          ;
        }

        MI x = dp[i][{j, l}];
        dp[i + 1][{j, l}] += x;
        dp[i + 1][{j, l + 2}] -= x;
        int w = min(j, v[i]);
        dp[i + 1][{j - w - 1, l + w + 1}] -= x;
        dp[i + 1][{j - w, l + w + 2}] += x;
        w = sum - j;
        dp[i + 1][{j + w + 1, l + w + 1}] -= x;
        dp[i + 1][{j + w, l + w + 2}] += x;
      }
    }
    /*
                    incII(l, 0, 2 * sum) {
                    incII(j, -1, sum + 1) {
                            cout << dp[i][{ j, l }] << " ";
                    } cout << endl;
                    } cout << endl;
    */
  }
  MI ans = 0;
  incII(l, 0, 2 * sum) {
    if (l <= 2 * k) {
      ans += dp[n][{0, l}];
    }
  }
  OUT(ans);

  /*
  OUT("---------");
  {
          auto dp = MV<MI>(0, n + 1, sum + 1, 2 * sum + 1);
          dp[0][0][0] = 1;
          inc(i, n) {
          incII(l, 0, 2 * sum) {
          incII(j, 0, sum) {
                  incII(w, 0, 2*sum) {
                          int jj = j + w - v[i];
                          int ll = l + abs(w - v[i]);
                          if(inII(jj, 0, sum) && inII(ll, 0, 2 * sum)) {
                                  dp[i + 1][jj][ll] += dp[i][j][l];
                          }
                  }
          }
          }
          }


                  inc(i, n + 1) {
                  incII(l, 0, 2 * sum) {
                  incII(j, 0, sum) {
                          cout << dp[i][j][l] << " ";
                  } cout << endl;
                  } cout << endl;
                  }
          MI ans = 0;
          incII(l, 0, 2 * sum) {
                  if(l <= 2 * k) { ans += dp[n][0][l]; }
          }
          OUT(ans);
  }
  */
}