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

#include <functional>
#include <iostream>
#include <numeric>
#include <vector>

template <class T> std::vector<T> vec(int len, T elem) {
  return std::vector<T>(len, elem);
}

template <int MOD> struct ModInt {
  using lint = long long;
  int val;

  // constructor
  ModInt(lint v = 0) : val(v % MOD) {
    if (val < 0)
      val += MOD;
  };

  // unary operator
  ModInt operator+() const { return ModInt(val); }
  ModInt operator-() const { return ModInt(MOD - val); }
  ModInt inv() const { return this->pow(MOD - 2); }

  // arithmetic
  ModInt operator+(const ModInt &x) const { return ModInt(*this) += x; }
  ModInt operator-(const ModInt &x) const { return ModInt(*this) -= x; }
  ModInt operator*(const ModInt &x) const { return ModInt(*this) *= x; }
  ModInt operator/(const ModInt &x) const { return ModInt(*this) /= x; }
  ModInt pow(lint n) const {
    auto x = ModInt(1);
    auto b = *this;
    while (n > 0) {
      if (n & 1)
        x *= b;
      n >>= 1;
      b *= b;
    }
    return x;
  }

  // compound assignment
  ModInt &operator+=(const ModInt &x) {
    if ((val += x.val) >= MOD)
      val -= MOD;
    return *this;
  }
  ModInt &operator-=(const ModInt &x) {
    if ((val -= x.val) < 0)
      val += MOD;
    return *this;
  }
  ModInt &operator*=(const ModInt &x) {
    val = lint(val) * x.val % MOD;
    return *this;
  }
  ModInt &operator/=(const ModInt &x) { return *this *= x.inv(); }

  // compare
  bool operator==(const ModInt &b) const { return val == b.val; }
  bool operator!=(const ModInt &b) const { return val != b.val; }
  bool operator<(const ModInt &b) const { return val < b.val; }
  bool operator<=(const ModInt &b) const { return val <= b.val; }
  bool operator>(const ModInt &b) const { return val > b.val; }
  bool operator>=(const ModInt &b) const { return val >= b.val; }

  // I/O
  friend std::istream &operator>>(std::istream &is, ModInt &x) noexcept {
    lint v;
    is >> v;
    x = v;
    return is;
  }
  friend std::ostream &operator<<(std::ostream &os, const ModInt &x) noexcept {
    return os << x.val;
  }
};

template <class T> struct SegmentTree {
  using Merger = std::function<T(T, T)>;

  int length;
  std::vector<T> dat;
  T unit;
  Merger merge;

  SegmentTree() = default;

  SegmentTree(int n, const T &unit, const Merger &merge)
      : length(1), unit(unit), merge(merge) {
    while (length < n)
      length <<= 1;
    dat.assign(length * 2, unit);
  }

  template <class Container>
  SegmentTree(const Container &elems, const T &unit, const Merger &merge)
      : length(1), unit(unit), merge(merge) {
    int n = elems.size();
    while (length < n)
      length <<= 1;

    dat.assign(length * 2, unit);

    std::copy(elems.begin(), elems.end(), dat.begin() + length);

    for (int nidx = length - 1; nidx >= 1; --nidx) {
      T vl = dat[nidx * 2 + 0];
      T vr = dat[nidx * 2 + 1];
      dat[nidx] = merge(vl, vr);
    }
  }

  void update(int nidx, const T &elem) {
    nidx += length;
    dat[nidx] = elem;

    while (nidx > 0) {
      nidx >>= 1;
      T vl = dat[nidx * 2 + 0];
      T vr = dat[nidx * 2 + 1];
      dat[nidx] = merge(vl, vr);
    }
  }

  T fold(int ql, int qr) const {
    ql = std::max(ql, 0);
    qr = std::min(qr, length);
    ql += length, qr += length;

    T lacc = unit, racc = unit;
    while (ql < qr) {
      if (ql & 1) {
        lacc = merge(lacc, dat[ql]);
        ++ql;
      }
      if (qr & 1) {
        --qr;
        racc = merge(dat[qr], racc);
      }
      ql >>= 1, qr >>= 1;
    }

    return merge(lacc, racc);
  }

  T get(int idx) const { return dat[idx + length]; }
  T fold_all() const { return dat[1]; }
};

constexpr int MOD = 1000000007;
using mint = ModInt<MOD>;

void solve() {
  int n, a, b;
  std::cin >> n >> a >> b;

  if (a > b)
    std::swap(a, b);
  // a <= b

  // the previous digit
  // index
  // the number of continuous 1s
  auto dp = vec(2, vec(n + 1, vec(b, mint(0))));
  dp[1][0][0] = 1;

  // the previous digit is 1
  SegmentTree<mint> seg(n + 1, 0, [](auto a, auto b) { return a + b; });
  seg.update(0, 1);

  for (int i = 1; i <= n; ++i) {
    for (int o = 1; o < b; ++o) {
      // put 1
      for (int zp = 0; zp < 2; ++zp) {
        dp[1][i][o] += dp[zp][i - 1][o - 1];
      }
    }

    // put short 0s after 1
    dp[0][i][0] = seg.fold(i - a + 1, i);

    // put a 0s after 1
    if (i >= a) {
      for (int po = 0; po + a < b; ++po) {
        dp[0][i][po + a] += dp[1][i - a][po];
      }
    }

    for (int o = a + 1; o < b; ++o) {
      // put 0 after long 0s
      dp[0][i][o] += dp[0][i - 1][o - 1];
    }

    mint sum = 0;
    for (auto x : dp[1][i])
      sum += x;
    seg.update(i, sum);
  }

  mint out = 0;
  for (int zp = 0; zp < 2; ++zp) {
    for (auto x : dp[zp][n])
      out += x;
  }

  std::cout << mint(2).pow(n) - out << "\\n";
}

int main() {
  std::cin.tie(nullptr);
  std::ios::sync_with_stdio(false);

  solve();

  return 0;
}


Error: TLE

Fixed submission:

#include <iostream>
#include <numeric>
#include <vector>

template <class T> std::vector<T> vec(int len, T elem) {
  return std::vector<T>(len, elem);
}

template <int MOD> struct ModInt {
  using lint = long long;
  int val;

  // constructor
  ModInt(lint v = 0) : val(v % MOD) {
    if (val < 0)
      val += MOD;
  };

  // unary operator
  ModInt operator+() const { return ModInt(val); }
  ModInt operator-() const { return ModInt(MOD - val); }
  ModInt inv() const { return this->pow(MOD - 2); }

  // arithmetic
  ModInt operator+(const ModInt &x) const { return ModInt(*this) += x; }
  ModInt operator-(const ModInt &x) const { return ModInt(*this) -= x; }
  ModInt operator*(const ModInt &x) const { return ModInt(*this) *= x; }
  ModInt operator/(const ModInt &x) const { return ModInt(*this) /= x; }
  ModInt pow(lint n) const {
    auto x = ModInt(1);
    auto b = *this;
    while (n > 0) {
      if (n & 1)
        x *= b;
      n >>= 1;
      b *= b;
    }
    return x;
  }

  // compound assignment
  ModInt &operator+=(const ModInt &x) {
    if ((val += x.val) >= MOD)
      val -= MOD;
    return *this;
  }
  ModInt &operator-=(const ModInt &x) {
    if ((val -= x.val) < 0)
      val += MOD;
    return *this;
  }
  ModInt &operator*=(const ModInt &x) {
    val = lint(val) * x.val % MOD;
    return *this;
  }
  ModInt &operator/=(const ModInt &x) { return *this *= x.inv(); }

  // compare
  bool operator==(const ModInt &b) const { return val == b.val; }
  bool operator!=(const ModInt &b) const { return val != b.val; }
  bool operator<(const ModInt &b) const { return val < b.val; }
  bool operator<=(const ModInt &b) const { return val <= b.val; }
  bool operator>(const ModInt &b) const { return val > b.val; }
  bool operator>=(const ModInt &b) const { return val >= b.val; }

  // I/O
  friend std::istream &operator>>(std::istream &is, ModInt &x) noexcept {
    lint v;
    is >> v;
    x = v;
    return is;
  }
  friend std::ostream &operator<<(std::ostream &os, const ModInt &x) noexcept {
    return os << x.val;
  }
};

constexpr int MOD = 1000000007;
using mint = ModInt<MOD>;

void solve() {
  int n, a, b;
  std::cin >> n >> a >> b;

  if (a > b)
    std::swap(a, b);
  // a <= b

  // the previous digit is 0?
  // index
  // the number of continuous 1s
  auto dp = vec(2, vec(n + 1, vec(b, mint(0))));
  dp[0][0][0] = 1;

  for (int i = 0; i < n; ++i) {
    for (int o = 0; o < b; ++o) {
      // don't put any 0s
      for (int zp = 0; zp < 2; ++zp) {
        if (o + 1 < b)
          dp[0][i + 1][o + 1] += dp[zp][i][o];
      }

      // put z 0s
      for (int z = 1; i + z <= n; ++z) {
        int no = (z < a ? 0 : o + z);
        if (no >= b)
          continue;
        dp[1][i + z][no] += dp[0][i][o];
      }
    }
  }

  mint out = 0;
  for (int zp = 0; zp < 2; ++zp) {
    for (int o = 0; o < b; ++o) {
      out += dp[zp][n][o];
    }
  }

  std::cout << mint(2).pow(n) - out << "\\n";
}

int main() {
  std::cin.tie(nullptr);
  std::ios::sync_with_stdio(false);

  solve();

  return 0;
}