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

#include <bits/stdc++.h>
typedef long long int ll;
typedef unsigned long long int ull;
#define BIG_NUM 2000000000
#define HUGE_NUM 99999999999999999
#define MOD 1000000007
#define EPS 0.000000001
using namespace std;

#define NUM 4001

struct Edge {
  Edge(short arg_to, ll arg_cost) {
    to = arg_to;
    cost = arg_cost;
  }
  bool operator<(const struct Edge &arg) const { return cost < arg.cost; }
  short to;
  ll cost;
};

struct Info {
  void set(int arg_loc, int arg_rank) {
    loc = arg_loc;
    rank = arg_rank;
  }
  short loc, rank;
};

int N;
Info pre[NUM][NUM];
ll table[NUM];
vector<short> MAX_DIVIDE[NUM]; // dp[場所][ある場所から出発するエッジ集合の中でのrank]
                               // = 到達までの最大分割数
vector<Edge> G[NUM];

int main() {

  scanf("%d", &N);

  table[0] = 0;

  for (int i = 1; i <= N; i++) {
    scanf("%lld", &table[i]);
    table[i] += table[i - 1];
  }

  // 各場所からエッジを張る
  for (int left = 0; left <= N - 1; left++) {
    for (int right = left + 1; right <= N; right++) {
      G[left].push_back(Edge(right, table[right] - table[left]));
    }
    sort(G[left].begin(), G[left].end());
  }

  G[N].push_back(Edge(N + 1, HUGE_NUM));

  int L, R, mid;
  int base_rank, start_index, next_loc, next_rank, out_rank;
  ll base_cost;

  // MAX_DIVIDE[場所][相対入次rank] = 左の最大分割数
  for (int loc = 0; loc <= N; loc++) {
    MAX_DIVIDE[loc].assign(N - loc + 1, -1);
  }

  int work[NUM];
  MAX_DIVIDE[0][0] = 0;

  for (int loc = 0; loc <= N - 1; loc++) {

    // ★本来の到着地点★
    for (int i = 0; i < G[loc].size(); i++) {
      work[i] = i;
    }

    // ある出発エッジについて、出発するまでの、左方最大分割数を求める
    for (int i = 1; i < G[loc].size(); i++) {
      if (MAX_DIVIDE[loc][i] < MAX_DIVIDE[loc][i - 1]) {
        MAX_DIVIDE[loc][i] = MAX_DIVIDE[loc][i - 1];
        work[i] = work[i - 1];
      }
    }

    for (int i = 0; i < G[loc].size(); i++) {

      if (MAX_DIVIDE[loc][i] == -1)
        continue;

      next_loc = G[loc][i].to;
      base_cost = G[loc][i].cost;

      // base_costよりも大きい、next_locにおける最小のエッジのインデックスを求める
      L = 0, R = G[next_loc].size() - 1, mid = (L + R) / 2;
      next_rank = BIG_NUM;

      while (L <= R) {
        if (G[next_loc][mid].cost > base_cost) {
          next_rank = mid;
          R = mid - 1;
        } else {
          L = mid + 1;
        }
        mid = (L + R) / 2;
      }

      if (next_rank == BIG_NUM)
        continue;

      if (MAX_DIVIDE[next_loc][next_rank] >= MAX_DIVIDE[loc][i] + 1)
        continue;

      MAX_DIVIDE[next_loc][next_rank] = MAX_DIVIDE[loc][i] + 1;
      pre[next_loc][next_rank].set(loc, work[i]);
    }
  }

  int maximum = MAX_DIVIDE[N][0], final_rank = 0;

  if (maximum == 1) {
    printf("1\\n\\n");
    return 0;
  }

  stack<int> S;

  int tmp_loc = N, tmp_rank = final_rank;
  Info info;

  while (pre[tmp_loc][tmp_rank].loc > 0) {
    S.push(pre[tmp_loc][tmp_rank].loc);
    info = pre[tmp_loc][tmp_rank];
    tmp_loc = info.loc;
    tmp_rank = info.rank;
  }

  printf("%d\\n", maximum);

  if (!S.empty()) {
    printf("%d", S.top());
    S.pop();

    while (!S.empty()) {
      printf(" %d", S.top());
      S.pop();
    }
  }
  printf("\\n");

  return 0;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>
typedef long long int ll;
typedef unsigned long long int ull;
#define BIG_NUM 2000000000
#define HUGE_NUM 99999999999999999
#define MOD 1000000007
#define EPS 0.000000001
using namespace std;

#define NUM 4001

struct Edge {
  Edge(short arg_to, ll arg_cost) {
    to = arg_to;
    cost = arg_cost;
  }
  bool operator<(const struct Edge &arg) const { return cost < arg.cost; }
  short to;
  ll cost;
};

struct Info {
  void set(int arg_loc, int arg_rank) {
    loc = arg_loc;
    rank = arg_rank;
  }
  short loc, rank;
};

int N;
Info pre[NUM][NUM];
ll table[NUM];
short MAX_DIVIDE[NUM][NUM]; // dp[場所][ある場所から出発するエッジ集合の中でのrank]
                            // = 到達までの最大分割数
vector<Edge> G[NUM];

int main() {

  scanf("%d", &N);

  table[0] = 0;

  for (int i = 1; i <= N; i++) {
    scanf("%lld", &table[i]);
    table[i] += table[i - 1];
  }

  // 各場所からエッジを張る
  for (int left = 0; left <= N - 1; left++) {
    for (int right = left + 1; right <= N; right++) {
      G[left].push_back(Edge(right, table[right] - table[left]));
    }
    sort(G[left].begin(), G[left].end());
  }

  G[N].push_back(Edge(N + 1, HUGE_NUM));

  int L, R, mid;
  int base_rank, start_index, next_loc, next_rank, out_rank;
  ll base_cost;

  // MAX_DIVIDE[場所][相対入次rank] = 左の最大分割数
  for (int loc = 0; loc <= N; loc++) {
    for (out_rank = 0; out_rank <= N - loc; out_rank++) { // その
      MAX_DIVIDE[loc][out_rank] = 0;
    }
  }

  int work[NUM];

  for (int loc = 0; loc <= N - 1; loc++) {

    // ★本来の到着地点★
    for (int i = 0; i < G[loc].size(); i++) {
      work[i] = i;
    }

    // ある出発エッジについて、出発するまでの、左方最大分割数を求める
    for (int i = 1; i < G[loc].size(); i++) {
      if (MAX_DIVIDE[loc][i] < MAX_DIVIDE[loc][i - 1]) {
        MAX_DIVIDE[loc][i] = MAX_DIVIDE[loc][i - 1];
        work[i] = work[i - 1];
      }
    }

    for (int i = 0; i < G[loc].size(); i++) {
      next_loc = G[loc][i].to;
      base_cost = G[loc][i].cost;

      // base_costよりも大きい、next_locにおける最小のエッジのインデックスを求める
      L = 0, R = G[next_loc].size() - 1, mid = (L + R) / 2;
      next_rank = BIG_NUM;

      while (L <= R) {
        if (G[next_loc][mid].cost > base_cost) {
          next_rank = mid;
          R = mid - 1;
        } else {
          L = mid + 1;
        }
        mid = (L + R) / 2;
      }

      if (next_rank == BIG_NUM)
        continue;

      if (MAX_DIVIDE[next_loc][next_rank] >= MAX_DIVIDE[loc][i] + 1)
        continue;

      MAX_DIVIDE[next_loc][next_rank] = MAX_DIVIDE[loc][i] + 1;
      pre[next_loc][next_rank].set(loc, work[i]);
    }
  }

  int maximum = MAX_DIVIDE[N][0], final_rank = 0;

  if (maximum == 1) {
    printf("1\\n\\n");
    return 0;
  }

  stack<int> S;

  int tmp_loc = N, tmp_rank = final_rank;
  Info info;

  while (pre[tmp_loc][tmp_rank].loc > 0) {
    S.push(pre[tmp_loc][tmp_rank].loc);
    info = pre[tmp_loc][tmp_rank];
    tmp_loc = info.loc;
    tmp_rank = info.rank;
  }

  printf("%d\\n", maximum);

  if (!S.empty()) {
    printf("%d", S.top());
    S.pop();

    while (!S.empty()) {
      printf(" %d", S.top());
      S.pop();
    }
  }
  printf("\\n");

  return 0;
}