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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
using namespace std;

typedef pair<int, int> P;
#define rep(i, n) for (int i = 0; i < (n); i++)
#define all(c) (c).begin(), (c).end()
#define uniq(c) c.erase(unique(all(c)), (c).end())
#define index(xs, x) (int)(lower_bound(all(xs), x) - xs.begin())
#define _1 first
#define _2 second
#define pb push_back
#define INF 10000
#define MOD 1000000007
inline int mp(int seg, int num) { return seg * 32767 + num; };
struct SegTree {
  int MAX_N;
  vector<int> seg;
  SegTree(int n) {
    MAX_N = 1;
    while (n >= MAX_N)
      MAX_N <<= 1;
    seg.resize(2 * MAX_N - 1, mp(-INF, -1));
  }

  int query(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r == -1)
      r = MAX_N;
    if (r <= a || b <= l)
      return mp(-INF, -1);
    if (a <= l && r <= b)
      return seg[k];
    return max(query(a, b, k * 2 + 1, l, (l + r) / 2),
               query(a, b, k * 2 + 2, (l + r) / 2, r));
  }

  void update(int k, P v) {
    k += MAX_N - 1;
    seg[k] = mp(v._1, v._2);
    while (k) {
      k = (k - 1) / 2;
      seg[k] = max(seg[k * 2 + 1], seg[k * 2 + 2]);
    }
  }
};

int N;
long long X[4001];
short pre[4001][4001];
vector<long long> ks[4001];

inline long long key(int i, int j) { return X[j] - X[i]; }
signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cin >> N;
  rep(i, N) cin >> X[i + 1];
  for (int i = 0; i < N; i++)
    X[i + 1] += X[i];

  vector<SegTree> seg;
  rep(i, N + 1) {
    rep(j, i) ks[i].pb(key(j, i));
    sort(all(ks[i]));
    uniq(ks[i]);
    seg.pb(SegTree(ks[i].size()));
  }

  for (int r = 1; r <= N; r++) {
    pre[0][r] = -1;
    seg[r].update(index(ks[r], key(0, r)), P(1, 0));
  }
  for (int l = 1; l <= N; l++) {
    for (int r = l + 1; r <= N; r++) {
      int p = seg[l].query(0, index(ks[l], key(l, r)));
      int p2 = p % 32767, p1 = p / 32767;
      if (p2 == -1)
        continue;
      pre[l][r] = p2;
      seg[r].update(index(ks[r], key(l, r)), P(p1 + 1, l));
    }
  }
  int m = seg[N].query(0, seg[N].MAX_N);
  int m2 = m % 32767, m1 = m / 32767;
  vector<int> ans;
  int left = m2, right = N;
  while (left != 0) {
    ans.pb(left);
    int k = left;
    left = pre[left][right];
    right = k;
  }
  reverse(all(ans));
  assert(ans.size() == m1 - 1);
  cout << m1 << "\\n";
  for (int x : ans)
    cout << x << " ";
  cout << "\\n";
  return 0;
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
using namespace std;

typedef pair<int, int> P;
#define rep(i, n) for (int i = 0; i < (n); i++)
#define all(c) (c).begin(), (c).end()
#define uniq(c) c.erase(unique(all(c)), (c).end())
#define index(xs, x) (int)(lower_bound(all(xs), x) - xs.begin())
#define _1 first
#define _2 second
#define pb push_back
#define INF 10000
#define MOD 1000000007
#define MAX_N (1 << 12)
inline int mp(int seg, int num) { return seg * 32767 + num; };
struct SegTree {
  int seg[2 * MAX_N - 1];
  SegTree() { rep(i, MAX_N * 2 - 1) seg[i] = mp(-INF, -1); }

  int query(int a, int b, int k = 0, int l = 0, int r = MAX_N) {
    if (r <= a || b <= l)
      return mp(-INF, -1);
    if (a <= l && r <= b)
      return seg[k];
    return max(query(a, b, k * 2 + 1, l, (l + r) / 2),
               query(a, b, k * 2 + 2, (l + r) / 2, r));
  }

  void update(int k, P v) {
    k += MAX_N - 1;
    seg[k] = mp(v._1, v._2);
    while (k) {
      k = (k - 1) / 2;
      seg[k] = max(seg[k * 2 + 1], seg[k * 2 + 2]);
    }
  }
};

SegTree seg[4001];

int N;
long long X[4001];
short pre[4001][4001];
vector<long long> ks[4001];

inline long long key(int i, int j) { return X[j] - X[i]; }
signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cin >> N;
  rep(i, N) cin >> X[i + 1];
  for (int i = 0; i < N; i++)
    X[i + 1] += X[i];

  rep(i, N + 1) {
    rep(j, i) ks[i].pb(key(j, i));
    sort(all(ks[i]));
    uniq(ks[i]);
  }

  for (int r = 1; r <= N; r++) {
    pre[0][r] = -1;
    seg[r].update(index(ks[r], key(0, r)), P(1, 0));
  }
  for (int l = 1; l <= N; l++) {
    for (int r = l + 1; r <= N; r++) {
      int p = seg[l].query(0, index(ks[l], key(l, r)));
      int p2 = p % 32767, p1 = p / 32767;
      if (p2 == -1)
        continue;
      pre[l][r] = p2;
      seg[r].update(index(ks[r], key(l, r)), P(p1 + 1, l));
    }
  }
  int m = seg[N].query(0, MAX_N);
  int m2 = m % 32767, m1 = m / 32767;
  vector<int> ans;
  int left = m2, right = N;
  while (left != 0) {
    ans.pb(left);
    int k = left;
    left = pre[left][right];
    right = k;
  }
  reverse(all(ans));
  assert(ans.size() == m1 - 1);
  cout << m1 << "\\n";
  for (int x : ans)
    cout << x << " ";
  cout << "\\n";
  return 0;
}