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

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

typedef vector<vector<int>> Graph;

typedef long long lli;

const lli MAXN = 500001;
const lli INF = 1LL << 60;

lli N, V;
Graph G;

int num[MAXN], depth[MAXN];
lli cost[MAXN], mem[MAXN];

/*
pair<lli, lli> dfs(lli v, lli d) {
  depth[v] = d;
  if(v >= N-1) {
    num[v] = 1;
    cost[v] = d;
    return make_pair(num[v], cost[v]);
  }
  num[v] = 0;
  cost[v] = 0;
  for(lli i = 0; i < G[v].size(); ++i) {
    pair<lli, lli> p = dfs(G[v][i], d + 1);
    num[v] += p.first;
    cost[v] += p.second;
  }
  return make_pair(num[v], cost[v]);
}

lli rec(lli v) {
  lli &res = mem[v];
  if(res != -1) return res;
  if(v >= N-1) return res = 0;
  res = INF;
  lli a = G[v][0];
  lli b = G[v][1];
  lli smallA = cost[a] - num[a] * depth[v];
  lli smallB = cost[b] - num[b] * depth[v];
  res = min(res, smallA + rec(b) + 1);
  res = min(res, rec(a) + smallB + 1);
  res = min(res, rec(a) + rec(b) + depth[v] + 2);
  return res;
}
*/

int main() {
  while (cin >> N) {
    V = N - 1;
    G = Graph(V);
    for (lli i = 0; i < N - 1; ++i) {
      lli a, b;
      cin >> a >> b;
      --a;
      --b;
      G[i].push_back(a);
      G[i].push_back(b);
    }

    {
      queue<lli> que;
      que.push(0);
      depth[0] = 0;
      while (que.size()) {
        lli v = que.front();
        que.pop();
        for (lli i = 0; i < G[v].size(); ++i) {
          lli nv = G[v][i];
          if (nv == N - 1)
            continue;
          que.push(nv);
          depth[nv] = depth[v] + 1;
        }
      }
    }

    vector<pair<lli, lli>> w;
    for (lli i = 0; i < V; ++i) {
      w.push_back(make_pair(-depth[i], i));
    }
    sort(w.begin(), w.end());

    for (lli i = 0; i < V; ++i) {
      lli v = w[i].second;
      if (G[v].empty()) {
        assert(false);
        cost[v] = depth[v];
        num[v] = 1;
      } else {
        lli a = G[v][0];
        lli b = G[v][1];
        cost[v] = num[v] = 0;
        if (a == N - 1) {
          cost[v] += depth[v] + 1;
          num[v] += 1;
        } else {
          cost[v] += cost[a];
          num[v] += num[a];
        }
        if (b == N - 1) {
          cost[v] += depth[v] + 1;
          num[v] += 1;
        } else {
          cost[v] += cost[b];
          num[v] += num[b];
        }
      }
    }

    mem[N - 1] = 0;
    for (lli i = 0; i < V; ++i) {
      lli v = w[i].second;
      lli &res = mem[v];
      res = INF;
      lli a = G[v][0];
      lli b = G[v][1];
      lli smallA = 0, smallB = 0;
      if (a == N - 1) {
        smallA = depth[v] + 1 - 1 * depth[v];
      } else {
        smallA = cost[a] - num[a] * depth[v];
      }

      if (b == N - 1) {
        smallB = depth[v] + 1 - 1 * depth[v];
      } else {
        smallB = cost[b] - num[b] * depth[v];
      }
      res = min(res, smallA + mem[b] + 1);
      res = min(res, mem[a] + smallB + 1);
      res = min(res, mem[a] + mem[b] + depth[v] + 2);
    }

    cout << mem[0] << endl;
  }
  return 0;
}

Error: MLE

Fixed submission:

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

typedef vector<vector<int>> Graph;

const int MAXN = 1000000;
const int INF = 1 << 28;

int N, V;
Graph G;

int num[MAXN], cost[MAXN], depth[MAXN];
int mem[MAXN];
/*
pair<int, int> dfs(int v, int d) {
  depth[v] = d;
  if(v >= N-1) {
    num[v] = 1;
    cost[v] = d;
    return make_pair(num[v], cost[v]);
  }
  num[v] = 0;
  cost[v] = 0;
  for(int i = 0; i < G[v].size(); ++i) {
    pair<int, int> p = dfs(G[v][i], d + 1);
    num[v] += p.first;
    cost[v] += p.second;
  }
  return make_pair(num[v], cost[v]);
}

int rec(int v) {
  int &res = mem[v];
  if(res != -1) return res;
  if(v >= N-1) return res = 0;
  res = INF;
  int a = G[v][0];
  int b = G[v][1];
  int smallA = cost[a] - num[a] * depth[v];
  int smallB = cost[b] - num[b] * depth[v];
  res = min(res, smallA + rec(b) + 1);
  res = min(res, rec(a) + smallB + 1);
  res = min(res, rec(a) + rec(b) + depth[v] + 2);
  return res;
}
*/

int main() {
  while (cin >> N) {
    V = 2 * N - 1;
    G = Graph(V);
    int leaf = N - 1;
    for (int i = 0; i < N - 1; ++i) {
      int a, b;
      cin >> a >> b;
      --a;
      --b;
      if (a == N - 1)
        a = leaf++;
      if (b == N - 1)
        b = leaf++;
      G[i].push_back(a);
      G[i].push_back(b);
    }

    {
      queue<int> que;
      que.push(0);
      depth[0] = 0;
      while (que.size()) {
        int v = que.front();
        que.pop();
        for (int i = 0; i < G[v].size(); ++i) {
          int nv = G[v][i];
          que.push(nv);
          depth[nv] = depth[v] + 1;
        }
      }
    }

    vector<pair<int, int>> w;
    for (int i = 0; i < V; ++i) {
      w.push_back(make_pair(-depth[i], i));
    }
    sort(w.begin(), w.end());

    for (int i = 0; i < V; ++i) {
      int v = w[i].second;
      if (G[v].empty()) {
        cost[v] = depth[v];
        num[v] = 1;
      } else {
        int a = G[v][0];
        int b = G[v][1];
        cost[v] = cost[a] + cost[b];
        num[v] = num[a] + num[b];
      }
    }

    memset(mem, -1, sizeof(mem));
    for (int i = 0; i < V; ++i) {
      int v = w[i].second;
      int &res = mem[v];
      if (v >= N - 1) {
        res = 0;
        continue;
      }
      res = INF;
      int a = G[v][0];
      int b = G[v][1];
      int smallA = cost[a] - num[a] * depth[v];
      int smallB = cost[b] - num[b] * depth[v];
      res = min(res, smallA + mem[b] + 1);
      res = min(res, mem[a] + smallB + 1);
      res = min(res, mem[a] + mem[b] + depth[v] + 2);
    }

    cout << mem[0] << endl;
  }
  return 0;
}