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

#include <algorithm>
#include <cstdio>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;

typedef long long ll;
typedef pair<ll, int> P;

#define rep(i, x) for (int i = 0; i < x; i++)
#define rep1(i, x) for (int i = 1; i <= x; i++)
#define pb push_back
#define sor(vec) sort(vec.begin(), vec.end())
#define lb(vec, a) lower_bound(vec.begin(), vec.end(), a)
#define ub(vec, a) upper_bound(vec.begin(), vec.end(), a)
#define fr first
#define sc second

const ll INF = 1000000000000000000;

struct bit {
  int siz;
  P s[8200];
  void init() {
    siz = 1 << 13;
    rep1(i, siz) { s[i] = P(-INF, 0); }
  }
  void updata(int a, P x) {
    while (a <= siz) {
      if (s[a].fr < x.fr)
        s[a] = x;
      a += a & -a;
    }
  }
  P query(int a) {
    P ret = P(-INF, 0);
    if (a == 0)
      return ret;
    while (a > 0) {
      if (ret.fr < s[a].fr)
        ret = s[a];
      a -= a & -a;
    }
    return ret;
  }
} bit;

int main() {
  int n;
  ll a[4002];
  scanf("%d", &n);
  rep1(i, n) scanf("%lld", &a[i]);

  ll sum[4002];
  static P dp[4002][4002];
  sum[0] = 0;
  rep1(i, n) sum[i] = sum[i - 1] + a[i];
  rep1(i, n) dp[i][1] = P(sum[i], 0);
  for (int j = 1; j <= n; j++) {
    vector<ll> Z;
    for (int i = 1; i <= n; i++) {
      Z.pb(sum[i] + dp[i][j].fr);
    }
    sor(Z);
    bit.init();
    for (int i = j; i <= n; i++) {
      dp[i][j + 1] = bit.query(lb(Z, sum[i]) - Z.begin());
      if (dp[i][j + 1].fr == -INF)
        dp[i][j + 1].fr = INF;
      else
        dp[i][j + 1].fr = sum[i] - dp[i][j + 1].fr;
      bit.updata((lb(Z, sum[i] + dp[i][j].fr) - Z.begin()) + 1, P(sum[i], i));
    }
  }
  for (int j = n; j >= 1; j--) {
    if (dp[n][j].fr != INF) {
      vector<int> ans;
      int loc = n;
      while (1) {
        loc = dp[loc][j].sc;
        j--;
        if (loc == 0)
          break;
        ans.pb(loc);
      }
      printf("%d\\n", (int)ans.size() + 1);
      reverse(ans.begin(), ans.end());
      rep(i, ans.size()) {
        printf("%d%c", ans[i], (i == (int)ans.size() - 1) ? '\\n' : ' ');
      }
      if (ans.size() == 0)
        puts("");
      return 0;
    }
  }
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <cstdio>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;

typedef long long ll;
typedef pair<ll, int> P;

#define rep(i, x) for (int i = 0; i < x; i++)
#define rep1(i, x) for (int i = 1; i <= x; i++)
#define pb push_back
#define sor(vec) sort(vec.begin(), vec.end())
#define lb(vec, a) lower_bound(vec.begin(), vec.end(), a)
#define ub(vec, a) upper_bound(vec.begin(), vec.end(), a)
#define fr first
#define sc second

const ll INF = 1000000000000000000;

struct bit {
  int siz;
  P s[8200];
  void init() {
    siz = 1 << 13;
    rep1(i, siz) { s[i] = P(-INF, 0); }
  }
  void updata(int a, P x) {
    while (a <= siz) {
      if (s[a].fr < x.fr)
        s[a] = x;
      a += a & -a;
    }
  }
  P query(int a) {
    P ret = P(-INF, 0);
    if (a == 0)
      return ret;
    while (a > 0) {
      if (ret.fr < s[a].fr)
        ret = s[a];
      a -= a & -a;
    }
    return ret;
  }
} bit;

int main() {
  int n;
  ll a[4002];
  scanf("%d", &n);
  rep1(i, n) scanf("%lld", &a[i]);

  ll sum[4002];
  static P dp[4002][4002];
  sum[0] = 0;
  rep1(i, n) sum[i] = sum[i - 1] + a[i];
  rep1(i, n) dp[i][1] = P(sum[i], 0);
  for (int j = 1; j <= n; j++) {
    vector<ll> Z;
    for (int i = 1; i <= n; i++) {
      Z.pb(sum[i] + dp[i][j].fr);
    }
    sor(Z);
    bit.init();
    for (int i = 1; i <= n; i++) {
      dp[i][j + 1] = bit.query(lb(Z, sum[i]) - Z.begin());
      if (dp[i][j + 1].fr == -INF)
        dp[i][j + 1].fr = INF;
      else
        dp[i][j + 1].fr = sum[i] - dp[i][j + 1].fr;
      bit.updata((lb(Z, sum[i] + dp[i][j].fr) - Z.begin()) + 1, P(sum[i], i));
    }
  }
  for (int j = n; j >= 1; j--) {
    if (dp[n][j].fr != INF) {
      vector<int> ans;
      int loc = n;
      while (1) {
        loc = dp[loc][j].sc;
        j--;
        if (loc == 0)
          break;
        ans.pb(loc);
      }
      printf("%d\\n", (int)ans.size() + 1);
      reverse(ans.begin(), ans.end());
      rep(i, ans.size()) {
        printf("%d%c", ans[i], (i == (int)ans.size() - 1) ? '\\n' : ' ');
      }
      if (ans.size() == 0)
        puts("");
      return 0;
    }
  }
}