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

#include <bits/stdc++.h>

#define REP(i, s, n) for (int i = s; i < n; i++)
#define rep(i, n) REP(i, 0, n)
#define EPS (1e-10)
#define equals(a, b) (fabs((a) - (b)) < EPS)

using namespace std;

bool LT(double a, double b) { return !equals(a, b) && a < b; }
bool LTE(double a, double b) { return equals(a, b) || a < b; }

struct Point {
  double x, y, z;
  Point operator-(const Point &p) const {
    return (Point){x - p.x, y - p.y, z - p.z};
  }
  bool operator<(const Point &p) const {
    return ((!equals(z, p.z)) ? z > p.z
                              : ((!equals(y, p.y)) ? y < p.y : LT(x, p.x)));
  }
};

typedef pair<double, int> ii;

struct Edge {
  int to, cap;
  double cost;
  int rev;
  Edge(int to = 0, int cap = 0, double cost = 0, int rev = 0)
      : to(to), cap(cap), cost(cost), rev(rev) {}
};

const int MAX_V = 1100, IINF = INT_MAX;
int V;
vector<Edge> G[MAX_V];
double h[MAX_V], dist[MAX_V];
int prevv[MAX_V], preve[MAX_V];

inline void add_edge(int from, int to, int cap, double cost) {
  G[from].push_back(Edge(to, cap, cost, G[to].size()));
  G[to].push_back(Edge(from, 0, -cost, G[from].size() - 1));
}

double min_cost_flow(int s, int t) {
  double res = 0;
  while (1) {
    fill(dist, dist + V, IINF);
    dist[s] = 0;
    bool update = true;
    while (update) {
      update = false;
      rep(v, V) {
        if (dist[v] == IINF)
          continue;
        rep(i, (int)G[v].size()) {
          Edge &e = G[v][i];
          if (e.cap > 0 && !equals(dist[e.to], dist[v] + e.cost) &&
              dist[e.to] > dist[v] + e.cost) {
            // cout << "update dist[" << e.to << "] = dist[" << v << "] + " <<
            // e.cost << endl;
            dist[e.to] = dist[v] + e.cost;
            // cout << dist[e.to] << " = " << dist[v] << " + " << e.cost << endl
            // << endl;
            prevv[e.to] = v;
            preve[e.to] = i;
            update = true;
          }
        }
      }
    }

    if (equals(dist[t], IINF))
      break;

    int d = IINF;
    for (int v = t; v != s; v = prevv[v]) {
      d = min(d, G[prevv[v]][preve[v]].cap);
    }
    if (d == 0)
      break;
    res += d * dist[t];
    for (int v = t; v != s; v = prevv[v]) {
      Edge &e = G[prevv[v]][preve[v]];
      e.cap -= d;
      G[v][e.rev].cap += d;
    }
  }
  return res;
}

int N, K;
Point ps[100];

const double MAX = 10000;

double getDist(Point p) { return sqrt(p.x * p.x + p.y * p.y + p.z * p.z); }

void compute() {
  map<int, int> counter;
  bool failed = false;
  rep(i, N) counter[(int)ps[i].z]++;
  for (map<int, int>::iterator it = counter.begin(); it != counter.end();
       it++) {
    if (it->second > K) {
      failed = true;
      break;
    }
  }
  if (failed) {
    puts("-1");
    return;
  }
  rep(i, N * 4) G[i].clear();
  int source = N * 2;
  int sink = source + 2;
  V = sink + 1;
  add_edge(source, source + 1, K, 0);
  rep(i, N) {
    add_edge(source + 1, i, 1, 0);
    add_edge(i, N + i, 1, -MAX);
    add_edge(N + i, sink, 1, 0);
  }

  rep(i, N) rep(j, N) if (i != j && ps[i].z > ps[j].z) {
    add_edge(N + i, j, 1, getDist(ps[i] - ps[j]));
  }

  printf("%.10f\\n", min_cost_flow(source, sink) + N * MAX);
}

int main() {
  while (cin >> N >> K, N | K) {
    rep(i, N) cin >> ps[i].x >> ps[i].y >> ps[i].z;
    compute();
  }
  return 0;
}

Error: MLE

Fixed submission:

#include <bits/stdc++.h>

#define REP(i, s, n) for (int i = s; i < n; i++)
#define rep(i, n) REP(i, 0, n)
#define EPS (1e-10)
#define equals(a, b) (fabs((a) - (b)) < EPS)

using namespace std;

bool LT(double a, double b) { return !equals(a, b) && a < b; }
bool LTE(double a, double b) { return equals(a, b) || a < b; }

struct Point {
  float x, y, z;
  Point operator-(const Point &p) const {
    return (Point){x - p.x, y - p.y, z - p.z};
  }
  bool operator<(const Point &p) const {
    return ((!equals(z, p.z)) ? z > p.z
                              : ((!equals(y, p.y)) ? y < p.y : LT(x, p.x)));
  }
};

const int MAX_V = 110; ////////////////////////////////////////////////////////
const float DINF = 1e10;
int N, K;
Point ps[MAX_V];
float memo[MAX_V][MAX_V][MAX_V][MAX_V];
typedef unsigned long long ull;
const ull base = 1000000007ULL;
// unordered_map<ull,double> memo;
double mini;

float abs(Point p) { return sqrt(p.x * p.x + p.y * p.y + p.z * p.z); }

float rec(int cur, int *prev, int kept) {
  if (cur >= N)
    return 0;
  // rep(i,4) cout << prev[i] << " "; cout << endl;
  int st[4];
  rep(i, 4) st[i] = prev[i];
  sort(st, st + 4);
  /*
  ull key = 0ULL;
  rep(i,4) {
    key *= base;
    key += (ull)(st[i]+1);
  }
  if( memo.count(key) ) return memo[key];
  double ret = DINF;
  */
  float &ret = memo[st[0] + 1][st[1] + 1][st[2] + 1][st[3] + 1];
  if (!equals(ret, DINF))
    return ret;

  rep(i, 4) {
    if (i && prev[i - 1] == -1)
      break;
    int temp = prev[i];
    if (prev[i] == -1) {
      if (kept + 1 > K)
        continue;
      if (i && prev[i - 1] > cur)
        continue;
      prev[i] = cur;
      ret = min(ret, rec(cur + 1, prev, kept + 1));
    } else {
      if (LT(ps[cur].z, ps[prev[i]].z)) {
        if (i && prev[i - 1] > cur)
          continue;
        prev[i] = cur;
        ret = min(ret, rec(cur + 1, prev, kept) + abs(ps[cur] - ps[temp]));
      }
    }
    prev[i] = temp;
  }

  return ret;
  // return memo[key] = ret;
}

int main() {
  while (cin >> N >> K, N | K) {
    map<int, int> counter;
    bool failed = false;
    rep(i, N) {
      cin >> ps[i].x >> ps[i].y >> ps[i].z;
      int cnt = ++counter[ps[i].z];
      if (cnt > K)
        failed = true;
    }
    if (failed) {
      puts("-1");
      continue;
    }
    sort(ps, ps + N);
    rep(i, N + 2) rep(j, N + 2) rep(k, N + 2) rep(l, N + 2) memo[i][j][k][l] =
        DINF;
    // memo.clear();
    int temp[4] = {-1, -1, -1, -1};
    mini = DINF;
    mini = rec(0, temp, 0);
    if (equals(mini, DINF))
      puts("-1");
    else
      printf("%.10lf\\n", mini);
  }
  return 0;
}