Problem p01544 (C++) — buggy submission, verdict: Memory Limit Exceeded

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

typedef long long ll;
typedef unsigned long long ull;
typedef pair<ll, ll> P;

#define fi first
#define se second
#define repl(i, a, b) for (ll i = (ll)(a); i < (ll)(b); i++)
#define rep(i, n) repl(i, 0, n)
#define all(x) (x).begin(), (x).end()
#define dbg(x) cout << #x "=" << x << endl
#define mmax(x, y) (x > y ? x : y)
#define mmin(x, y) (x < y ? x : y)
#define maxch(x, y) x = mmax(x, y)
#define minch(x, y) x = mmin(x, y)
#define uni(x) x.erase(unique(all(x)), x.end())
#define exist(x, y) (find(all(x), y) != x.end())
#define bcnt __builtin_popcountll

#define INF 1e16
template <typename T, typename E> struct SegmentTree {
  typedef function<T(T, T)> F;
  typedef function<T(T, E)> G;
  int n;
  F f;
  G g;
  T d1;
  vector<T> dat;
  // 要素数、結合演算、作用素の適用演算、台集合の単位元、初期値
  SegmentTree(int n_, F f, G g, T d1, vector<T> v = vector<T>()) {
    this->f = f;
    this->g = g;
    this->d1 = d1;
    init(n_);
    if (n_ == (int)v.size())
      build(n_, v);
  }
  void init(int n_) {
    n = 1;
    while (n < n_)
      n *= 2;
    dat.clear();
    dat.resize(2 * n - 1, d1);
  }
  void build(int n_, vector<T> v) {
    for (int i = 0; i < n_; i++)
      dat[i + n - 1] = v[i];
    for (int i = n - 2; i >= 0; i--)
      dat[i] = f(dat[i * 2 + 1], dat[i * 2 + 2]);
  }
  void update(int k, E a) {
    k += n - 1;
    dat[k] = g(dat[k], a);
    while (k > 0) {
      k = (k - 1) / 2;
      dat[k] = f(dat[k * 2 + 1], dat[k * 2 + 2]);
    }
  }
  T query(int a, int b, int k, int l, int r) {
    if (r <= a || b <= l)
      return d1;
    if (a <= l && r <= b)
      return dat[k];
    T vl = query(a, b, k * 2 + 1, l, (l + r) / 2);
    T vr = query(a, b, k * 2 + 2, (l + r) / 2, r);
    return f(vl, vr);
  }
  T query(int a, int b) { return query(a, b, 0, 0, n); }
};

ll N;
ll A[4001];
ll dp[2][4001];
ll pre[4001][4001];
ll sum[4001];

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

  cin >> N;
  rep(i, N) cin >> A[i];
  rep(i, N) { sum[i + 1] = sum[i] + A[i]; }
  rep(i, 2) rep(j, N + 1) { dp[i][j] = INF; }
  rep(i, N + 1) rep(j, N + 1) { pre[i][j] = -1; }
  ll crt = 0, nxt = 1;
  dp[crt][0] = -INF;
  ll ans = 0;
  rep(i, N) {
    SegmentTree<P, P> sgt(
        N + 1, [](P a, P b) { return min(a, b); }, [](P a, P b) { return b; },
        P(INF, INF));
    vector<P> vs;
    rep(j, N + 1) {
      vs.push_back(P(dp[crt][j] + sum[j], j));
      dp[nxt][j] = INF;
    }
    sort(all(vs));
    rep(j, N + 1) {
      ll k = lower_bound(all(vs), P(sum[j], -INF)) - vs.begin();
      P p = sgt.query(0, k);
      if (p.fi != INF) {
        if (j == N) {
          ans = i + 1;
        }
        dp[nxt][j] = p.fi + sum[j];
        pre[i + 1][j] = p.se;
      }

      k = lower_bound(all(vs), P(dp[crt][j] + sum[j], j)) - vs.begin();
      sgt.update(k, P(-sum[j], j));
    }
    swap(crt, nxt);
  }
  P crtp = P(ans, N);
  vector<ll> res;
  while (P(crtp.fi - 1, pre[crtp.fi][crtp.se]) != P(-1, -1)) {
    res.push_back(crtp.se);
    crtp = P(crtp.fi - 1, pre[crtp.fi][crtp.se]);
  }
  reverse(all(res));
  cout << ans << endl;
  if (res.size() > 1) {
    cout << res[0];
    repl(i, 1, res.size() - 1) cout << " " << res[i];
    cout << endl;
  }

  return 0;
}


Error: MLE

Fixed submission:

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

typedef long long ll;
typedef unsigned long long ull;
typedef pair<ll, ll> P;

#define fi first
#define se second
#define repl(i, a, b) for (ll i = (ll)(a); i < (ll)(b); i++)
#define rep(i, n) repl(i, 0, n)
#define all(x) (x).begin(), (x).end()
#define dbg(x) cout << #x "=" << x << endl
#define mmax(x, y) (x > y ? x : y)
#define mmin(x, y) (x < y ? x : y)
#define maxch(x, y) x = mmax(x, y)
#define minch(x, y) x = mmin(x, y)
#define uni(x) x.erase(unique(all(x)), x.end())
#define exist(x, y) (find(all(x), y) != x.end())
#define bcnt __builtin_popcountll

#define INF 1e16
template <typename T, typename E> struct SegmentTree {
  typedef function<T(T, T)> F;
  typedef function<T(T, E)> G;
  int n;
  F f;
  G g;
  T d1;
  vector<T> dat;
  // 要素数、結合演算、作用素の適用演算、台集合の単位元、初期値
  SegmentTree(int n_, F f, G g, T d1, vector<T> v = vector<T>()) {
    this->f = f;
    this->g = g;
    this->d1 = d1;
    init(n_);
    if (n_ == (int)v.size())
      build(n_, v);
  }
  void init(int n_) {
    n = 1;
    while (n < n_)
      n *= 2;
    dat.clear();
    dat.resize(2 * n - 1, d1);
  }
  void build(int n_, vector<T> v) {
    for (int i = 0; i < n_; i++)
      dat[i + n - 1] = v[i];
    for (int i = n - 2; i >= 0; i--)
      dat[i] = f(dat[i * 2 + 1], dat[i * 2 + 2]);
  }
  void update(int k, E a) {
    k += n - 1;
    dat[k] = g(dat[k], a);
    while (k > 0) {
      k = (k - 1) / 2;
      dat[k] = f(dat[k * 2 + 1], dat[k * 2 + 2]);
    }
  }
  T query(int a, int b, int k, int l, int r) {
    if (r <= a || b <= l)
      return d1;
    if (a <= l && r <= b)
      return dat[k];
    T vl = query(a, b, k * 2 + 1, l, (l + r) / 2);
    T vr = query(a, b, k * 2 + 2, (l + r) / 2, r);
    return f(vl, vr);
  }
  T query(int a, int b) { return query(a, b, 0, 0, n); }
};

ll N;
ll A[4001];
vector<P> acc[4001];
ll dp[4001][4001];
ll pre[4001][4001];
ll sum[4001];
vector<ll> ord;

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

  cin >> N;
  rep(i, N) cin >> A[i];
  rep(i, N) {
    sum[i + 1] = sum[i] + A[i];
    ord.push_back(i + 1);
  }
  rep(i, N + 1) rep(j, N + 1) {
    dp[i][j] = INF;
    pre[i][j] = -1;
  }
  dp[0][0] = -INF;
  rep(i, N) {
    SegmentTree<P, P> sgt(
        N + 1, [](P a, P b) { return min(a, b); }, [](P a, P b) { return b; },
        P(INF, INF));
    vector<P> vs;
    rep(j, N + 1) { vs.push_back(P(dp[i][j] + sum[j], j)); }
    sort(all(vs));
    rep(j, N + 1) {
      ll k = lower_bound(all(vs), P(sum[j], -INF)) - vs.begin();
      P p = sgt.query(0, k);
      if (p.fi != INF) {
        dp[i + 1][j] = p.fi + sum[j];
        pre[i + 1][j] = p.se;
      }

      k = lower_bound(all(vs), P(dp[i][j] + sum[j], j)) - vs.begin();
      sgt.update(k, P(-sum[j], j));
    }
  }
  ll ans = 0;
  rep(j, N + 1) {
    if (dp[j][N] != INF)
      ans = j;
  }
  P crt = P(ans, N);
  vector<ll> res;
  while (P(crt.fi - 1, pre[crt.fi][crt.se]) != P(-1, -1)) {
    res.push_back(crt.se);
    crt = P(crt.fi - 1, pre[crt.fi][crt.se]);
  }
  reverse(all(res));
  cout << ans << endl;
  if (res.size() > 1) {
    cout << res[0];
    repl(i, 1, res.size() - 1) cout << " " << res[i];
    cout << endl;
  }

  return 0;
}