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

#include <algorithm>
#include <functional>
#include <inttypes.h>
#include <queue>
#include <stdint.h>
#include <stdio.h>

int ri() {
  int n;
  scanf("%d", &n);
  return n;
}

#define INF 1000000000000000000

struct Node;
extern Node *NONE;
struct Node {
  template <class T> struct Min {
    std::priority_queue<T, std::vector<T>, std::greater<T>> all;
    std::priority_queue<T, std::vector<T>, std::greater<T>> removing;
    void add(T content) { all.push(content); }
    void fix() {
      while (removing.size() && removing.top() == all.top())
        removing.pop(), all.pop();
    }
    void remove(T content) {
      removing.push(content);
      fix();
    }
    T min() { return all.size() ? all.top() : T{INF, {}}; }
    int size() { return all.size() - removing.size(); }
  };
#define l ch[0]
#define r ch[1]
  Node *ch[2] = {NONE, NONE};
  Node *p = NONE;
  /* vvvv   LIGHT stuff   vvvv */
  Min<std::pair<int64_t, int>> light_nearest[2];
  Min<std::pair<int64_t, std::pair<int, int>>> light_res;
  /* vvvv   self stuff   vvvv */
  bool rev = false;
  bool marked[2] = {false, false};
  int len = 0;
  int id = -1;
  int bias = 0;
  /* vvvv   heavy stuff   vvvv */
  int64_t len_sum = 0;
  std::pair<int64_t, std::pair<int, int>> res = {INF, {-1, -1}};
  std::pair<int64_t, int> left[2] = {{INF, -1}, {INF, -1}};
  std::pair<int64_t, int> right[2] = {{INF, -1}, {INF, -1}};
  Node *fetch() {
    if (l != NONE)
      l->flush();
    if (r != NONE)
      r->flush();
    len_sum = len + l->len_sum + r->len_sum;
    // a lot of candidates
    res = {INF, {-1, -1}};
    if (light_res.size())
      res = light_res.min(); // inside light
    res = std::min(
        res, {light_nearest[0].min().first + light_nearest[1].min().first + len,
              {light_nearest[0].min().second,
               light_nearest[1].min().second}}); // light to light
    if (marked[0]) {                             // from this
      res = std::min(res, {l->right[1].first + bias,
                           {id, l->right[1].second}}); // heavy left
      res = std::min(res, {r->left[1].first + bias,
                           {id, r->left[1].second}}); // heavy right
      res = std::min(res, {light_nearest[1].min().first + bias,
                           {id, light_nearest[1].min().second}}); // light
    }
    if (marked[1]) { // from this
      res = std::min(res, {l->right[0].first + bias,
                           {l->right[0].second, id}}); // heavy left
      res = std::min(res, {r->left[0].first + bias,
                           {r->left[0].second, id}}); // heavy right
      res = std::min(res, {light_nearest[0].min().first + bias,
                           {light_nearest[0].min().second, id}}); // light
    }
    // left to right
    res = std::min(res, {l->right[0].first + r->left[1].first + len,
                         {l->right[0].second, r->left[1].second}});
    res = std::min(res, {l->right[1].first + r->left[0].first + len,
                         {r->left[0].second, l->right[1].second}});
    // light to left/right
    res = std::min(res, {l->right[0].first + light_nearest[1].min().first + len,
                         {l->right[0].second, light_nearest[1].min().second}});
    res = std::min(res, {r->left[0].first + light_nearest[1].min().first + len,
                         {r->left[0].second, light_nearest[1].min().second}});
    res = std::min(res, {l->right[1].first + light_nearest[0].min().first + len,
                         {light_nearest[0].min().second, l->right[1].second}});
    res = std::min(res, {r->left[1].first + light_nearest[0].min().first + len,
                         {light_nearest[0].min().second, r->left[1].second}});
    res = std::min({res, l->res, r->res}); // inside heavy left/right

    for (int i = 0; i < 2; i++) {
      left[i] =
          std::min({l->left[i],
                    {r->left[i].first + len + l->len_sum, r->left[i].second},
                    {light_nearest[i].min().first + len + l->len_sum,
                     light_nearest[i].min().second}});
      if (marked[i])
        left[i] = std::min(left[i], {l->len_sum + bias, id});
      right[i] =
          std::min({r->right[i],
                    {l->right[i].first + len + r->len_sum, l->right[i].second},
                    {light_nearest[i].min().first + len + r->len_sum,
                     light_nearest[i].min().second}});
      if (marked[i])
        right[i] = std::min(right[i], {r->len_sum + bias, id});
    }
    return this;
  }
  Node *flush() {
    if (rev) {
      std::swap(l, r);
      l->rev ^= 1;
      r->rev ^= 1;
      std::swap(left[0], right[0]);
      std::swap(left[1], right[1]);
      rev = false;
    }
    return this;
  }
  Node *rotate(int dir) {
    Node *new_root = ch[!dir];
    ch[!dir] = new_root->ch[dir];
    ch[!dir]->p = this;
    new_root->ch[dir] = this;
    new_root->p = p;
    if (p->l == this)
      p->l = new_root;
    if (p->r == this)
      p->r = new_root;
    p = new_root;
    return fetch(), new_root->fetch();
  }
  bool is_root() { return p == NONE || (p->l != this && p->r != this); }
  void splay() {
    flush();
    while (!is_root()) {
      if (p->is_root()) {
        p->flush(), flush();
        p->rotate(p->l == this);
      } else {
        Node *pp = p->p;
        pp->flush(), p->flush(), flush();
        bool flag0 = pp->l == p;
        bool flag1 = p->l == this;
        if (flag0 == flag1)
          pp->rotate(flag0);
        p->rotate(flag1);
        if (flag0 != flag1)
          pp->rotate(flag0);
      }
    }
  }
  Node *expose() {
    Node *prev = NONE;
    for (Node *cur = this; cur != NONE; cur = cur->p) {
      cur->splay();
      if (cur->r != NONE) { // add
        cur->r->flush();
        for (int i = 0; i < 2; i++)
          cur->light_nearest[i].add(cur->r->left[i]);
        cur->light_res.add(cur->r->res);
      }
      cur->r = prev;
      if (cur->r != NONE) { // remove
        for (int i = 0; i < 2; i++)
          cur->light_nearest[i].remove(cur->r->left[i]);
        cur->light_res.remove(cur->r->res);
      }
      cur->fetch();
      prev = cur;
    }
    splay();
    return prev;
  }
  void link(Node *parent) {
    expose();
    parent->expose();
    p = parent;
    parent->r = this;
    parent->fetch();
  }
  void cut() {
    expose();
    l->p = NONE;
    l = NONE;
    fetch();
  }
  void evert() {
    expose();
    rev = true;
    flush();
  }
#undef l
#undef r
};
Node *NONE = new Node;

int main() {
  int n = ri();
  int a[n];
  for (auto &i : a)
    i = ri();
  Node nodes[n];
  Node edges[n - 1];
  for (int i = 0; i < n; i++) {
    nodes[i].id = i;
    nodes[i].marked[1] = true;
    nodes[i].bias = a[i];
    nodes[i].fetch();
  }
  for (int i = 0; i + 1 < n; i++) {
    int a = ri() - 1;
    int b = ri() - 1;
    int c = ri();
    edges[i].len = c;
    edges[i].fetch();
    nodes[a].evert();
    nodes[a].link(edges + i);
    edges[i].link(nodes + b);
  }
  auto flip = [&](int i) {
    nodes[i].expose();
    nodes[i].marked[1] = false;
    nodes[i].marked[0] = true;
    nodes[i].fetch();
  };
  flip(0);
  int64_t res = 0;
  for (int i = 0; i + 1 < n; i++) {
    nodes[0].evert();
    res += nodes[0].res.first;
    flip(nodes[0].res.second.second);
  }
  printf("%" PRId64 "\\n", res);
  return 0;
}


Error: MLE

Fixed submission:

#include <algorithm>
#include <functional>
#include <inttypes.h>
#include <queue>
#include <stdint.h>
#include <stdio.h>

int ri() {
  int n;
  scanf("%d", &n);
  return n;
}

#define INF 1000000000000000000

struct Node;
extern Node *NONE;
struct Node {
  template <class T> struct Min {
    std::priority_queue<T, std::vector<T>, std::greater<T>> all;
    std::priority_queue<T, std::vector<T>, std::greater<T>> removing;
    void add(T content) { all.push(content); }
    void fix() {
      while (removing.size() && removing.top() == all.top())
        removing.pop(), all.pop();
    }
    void remove(T content) {
      removing.push(content);
      fix();
    }
    T min() { return all.size() ? all.top() : T{INF, {}}; }
    int size() { return all.size() - removing.size(); }
  };
#define l ch[0]
#define r ch[1]
  Node *ch[2] = {NONE, NONE};
  Node *p = NONE;
  /* vvvv   LIGHT stuff   vvvv */
  Min<std::pair<int64_t, int>> light_nearest[2];
  Min<std::pair<int64_t, std::pair<int, int>>> light_res;
  /* vvvv   self stuff   vvvv */
  bool rev = false;
  bool marked[2] = {false, false};
  int len = 0;
  int id = -1;
  /* vvvv   heavy stuff   vvvv */
  int64_t len_sum = 0;
  std::pair<int64_t, std::pair<int, int>> res = {INF, {-1, -1}};
  std::pair<int64_t, int> left[2] = {{INF, -1}, {INF, -1}};
  std::pair<int64_t, int> right[2] = {{INF, -1}, {INF, -1}};
  Node *fetch() {
    if (l != NONE)
      l->flush();
    if (r != NONE)
      r->flush();
    len_sum = len + l->len_sum + r->len_sum;
    // a lot of candidates
    res = {INF, {-1, -1}};
    if (light_res.size())
      res = light_res.min(); // inside light
    res = std::min(
        res, {light_nearest[0].min().first + light_nearest[1].min().first + len,
              {light_nearest[0].min().second,
               light_nearest[1].min().second}}); // light to light
    if (marked[0]) {                             // from this
      res = std::min(
          res, {l->right[1].first, {id, l->right[1].second}}); // heavy left
      res = std::min(
          res, {r->left[1].first, {id, r->left[1].second}}); // heavy right
      res = std::min(res, {light_nearest[1].min().first,
                           {id, light_nearest[1].min().second}}); // light
    }
    if (marked[1]) { // from this
      res = std::min(
          res, {l->right[0].first, {l->right[0].second, id}}); // heavy left
      res = std::min(
          res, {r->left[0].first, {r->left[0].second, id}}); // heavy right
      res = std::min(res, {light_nearest[0].min().first,
                           {light_nearest[0].min().second, id}}); // light
    }
    // left to right
    res = std::min(res, {l->right[0].first + r->left[1].first + len,
                         {l->right[0].second, r->left[1].second}});
    res = std::min(res, {l->right[1].first + r->left[0].first + len,
                         {r->left[0].second, l->right[1].second}});
    // light to left/right
    res = std::min(res, {l->right[0].first + light_nearest[1].min().first + len,
                         {l->right[0].second, light_nearest[1].min().second}});
    res = std::min(res, {r->left[0].first + light_nearest[1].min().first + len,
                         {r->left[0].second, light_nearest[1].min().second}});
    res = std::min(res, {l->right[1].first + light_nearest[0].min().first + len,
                         {light_nearest[0].min().second, l->right[1].second}});
    res = std::min(res, {r->left[1].first + light_nearest[0].min().first + len,
                         {light_nearest[0].min().second, r->left[1].second}});
    res = std::min({res, l->res, r->res}); // inside heavy left/right

    for (int i = 0; i < 2; i++) {
      left[i] =
          std::min({l->left[i],
                    {r->left[i].first + len + l->len_sum, r->left[i].second},
                    {light_nearest[i].min().first + len + l->len_sum,
                     light_nearest[i].min().second}});
      if (marked[i])
        left[i] = std::min(left[i], {l->len_sum, id});
      right[i] =
          std::min({r->right[i],
                    {l->right[i].first + len + r->len_sum, l->right[i].second},
                    {light_nearest[i].min().first + len + r->len_sum,
                     light_nearest[i].min().second}});
      if (marked[i])
        right[i] = std::min(right[i], {r->len_sum, id});
    }
    return this;
  }
  Node *flush() {
    if (rev) {
      std::swap(l, r);
      l->rev ^= 1;
      r->rev ^= 1;
      std::swap(left[0], right[0]);
      std::swap(left[1], right[1]);
      rev = false;
    }
    return this;
  }
  Node *rotate(int dir) {
    Node *new_root = ch[!dir];
    ch[!dir] = new_root->ch[dir];
    ch[!dir]->p = this;
    new_root->ch[dir] = this;
    new_root->p = p;
    if (p->l == this)
      p->l = new_root;
    if (p->r == this)
      p->r = new_root;
    p = new_root;
    return fetch(), new_root->fetch();
  }
  bool is_root() { return p == NONE || (p->l != this && p->r != this); }
  void splay() {
    flush();
    while (!is_root()) {
      if (p->is_root()) {
        p->flush(), flush();
        p->rotate(p->l == this);
      } else {
        Node *pp = p->p;
        pp->flush(), p->flush(), flush();
        bool flag0 = pp->l == p;
        bool flag1 = p->l == this;
        if (flag0 == flag1)
          pp->rotate(flag0);
        p->rotate(flag1);
        if (flag0 != flag1)
          pp->rotate(flag0);
      }
    }
  }
  Node *expose() {
    Node *prev = NONE;
    for (Node *cur = this; cur != NONE; cur = cur->p) {
      cur->splay();
      if (cur->r != NONE) { // add
        cur->r->flush();
        for (int i = 0; i < 2; i++)
          cur->light_nearest[i].add(cur->r->left[i]);
        cur->light_res.add(cur->r->res);
      }
      cur->r = prev;
      if (cur->r != NONE) { // remove
        for (int i = 0; i < 2; i++)
          cur->light_nearest[i].remove(cur->r->left[i]);
        cur->light_res.remove(cur->r->res);
      }
      cur->fetch();
      prev = cur;
    }
    splay();
    return prev;
  }
  void link(Node *parent) {
    expose();
    parent->expose();
    p = parent;
    parent->r = this;
    parent->fetch();
  }
  void cut() {
    expose();
    l->p = NONE;
    l = NONE;
    fetch();
  }
  void evert() {
    expose();
    rev = true;
    flush();
  }
#undef l
#undef r
};
Node *NONE = new Node;

int main() {
  int n = ri();
  int a[n];
  for (auto &i : a)
    i = ri();
  Node nodes[2 * n];
  Node edges[2 * n - 1];
  for (int i = 0; i < 2 * n; i++) {
    nodes[i].id = i;
    if (i >= n)
      nodes[i].marked[1] = true;
    nodes[i].fetch();
  }
  for (int i = 0; i + 1 < n; i++) {
    int a = ri() - 1;
    int b = ri() - 1;
    int c = ri();
    edges[i].len = c;
    edges[i].fetch();
    nodes[a].evert();
    nodes[a].link(edges + i);
    edges[i].link(nodes + b);
  }
  for (int i = 0; i < n; i++) {
    edges[i + n - 1].len = a[i];
    edges[i + n - 1].fetch();
    edges[i + n - 1].link(nodes + i);
    nodes[i + n].link(edges + i + n - 1);
  }
  auto flip = [&](int i) {
    nodes[i].expose();
    nodes[i].marked[1] = false;
    nodes[i].marked[0] = true;
    nodes[i].fetch();
  };
  flip(n);
  int64_t res = 0;
  for (int i = 0; i + 1 < n; i++) {
    nodes[0].evert();
    res += nodes[0].res.first;
    flip(nodes[0].res.second.second);
  }
  printf("%" PRId64 "\\n", res);
  return 0;
}