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

#include <bits/stdc++.h>
#define foru(i, a, b) for (int i = int(a); i <= int(b); i++)
#define ford(i, a, b) for (int i = int(a); i >= int(b); i--)
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define sz(a) (int)a.size()
using namespace std;
typedef pair<int, int> pii;
typedef pair<int, pii> edge;

const int N = 200005;
const int oo = 1000000000;

int n, m, k, q, fes[N], dis[N];
vector<int> adj[N], wei[N];
edge e[2 * N];
/// tree + lca
int lab[N], p[N][25], fm[N][25], l[N];

void read() {
  scanf("%d%d%d%d", &n, &m, &k, &q);
  int u, v, w;
  foru(i, 1, m) {
    scanf("%d%d%d", &u, &v, &w);
    adj[u].pb(v);
    wei[u].pb(w);
    adj[v].pb(u);
    wei[v].pb(w);
  }
  foru(i, 1, k) {
    scanf("%d", &u);
    fes[u] = 1;
  }
}
/// min path
priority_queue<pii, vector<pii>, greater<pii>> heap;
void dijkstra() {
  foru(u, 1, n) if (fes[u]) {
    dis[u] = 0;
    heap.push(mp(0, u));
  }
  else dis[u] = oo;

  while (!heap.empty()) {
    int u = heap.top().se, e = heap.top().fi;
    heap.pop();
    if (dis[u] != e)
      continue;
    fes[u] = 1;
    foru(i, 0, sz(adj[u]) - 1) {
      int v = adj[u][i];
      if (fes[v])
        continue;
      if (dis[v] > dis[u] + wei[u][i]) {
        dis[v] = dis[u] + wei[u][i];
        heap.push(mp(dis[v], v));
      }
    }
  }
}
/// disjoint set
int root(int u) {
  if (lab[u] == u)
    return u;
  lab[u] = root(lab[u]);
  return lab[u];
}
void disjoint_set() {
  foru(u, 1, n) lab[u] = u;
  int sedge = 0;
  ford(i, m, 1) {
    int u = e[i].se.fi, v = e[i].se.se;
    u = root(u), v = root(v);
    if (u == v)
      continue;
    lab[u] = v;
    sedge++;
    adj[u].pb(v);
    wei[u].pb(e[i].fi);
    adj[v].pb(u);
    wei[v].pb(e[i].fi);
    if (sedge == n - 1)
      break;
  }
}
/// lca
void dfs(int u, int par) {
  foru(i, 0, sz(adj[u]) - 1) {
    int v = adj[u][i];
    if (par == v)
      continue;
    l[v] = l[u] + 1;
    fm[v][0] = wei[u][i];
    p[v][0] = u;
    dfs(v, u);
  }
}
int lca(int u, int v) {
  if (l[u] > l[v])
    swap(u, v);
  if (u == v)
    return 0;
  int res = oo;
  ford(j, 20, 0) if (p[v][j] > 0 && l[p[v][j]] >= l[u]) {
    res = min(res, fm[v][j]);
    v = p[v][j];
  }
  if (u == v)
    return res;
  ford(j, 20, 0) if (p[u][j] > 0 && p[v][j] > 0) if (p[u][j] != p[v][j]) {
    res = min(res, fm[u][j]);
    u = p[u][j];
    res = min(res, fm[v][j]);
    v = p[v][j];
  }
  res = min(res, min(fm[u][0], fm[v][0]));
  return res;
}
/// sol query

int main() {
  // freopen("a.txt", "r", stdin);
  read();
  dijkstra();
  m = 0;
  foru(u, 1, n) {
    foru(i, 0, sz(adj[u]) - 1) {
      int v = adj[u][i];
      e[++m] = mp(min(dis[u], dis[v]), mp(u, v));
    }
    adj[u].clear();
    wei[u].clear();
  }
  sort(e + 1, e + m + 1);
  disjoint_set();

  p[1][0] = -1;
  dfs(1, -1);
  foru(j, 1, 20) foru(u, 1, n) if (p[u][j - 1] > 0) {
    fm[u][j] = min(fm[u][j - 1], fm[p[u][j - 1]][j - 1]);
    p[u][j] = p[p[u][j - 1]][j - 1];
  }
  foru(i, 1, q) {
    int u, v;
    scanf("%d%d", &u, &v);
    printf("%d\\n", lca(u, v));
  }
  return 0;
}

Error: MLE

Fixed submission:

#include <bits/stdc++.h>
#define foru(i, a, b) for (int i = int(a); i <= int(b); i++)
#define ford(i, a, b) for (int i = int(a); i >= int(b); i--)
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define sz(a) (int)a.size()
using namespace std;
typedef pair<int, int> pii;
typedef pair<int, pii> edge;

const int N = 500005;
const int oo = 1000000000;

int n, m, k, q, fes[N], dis[N];
vector<int> adj[N], wei[N];
edge e[2 * N];
/// tree + lca
int lab[N], p[N][25], fm[N][25], l[N];

void read() {
  scanf("%d%d%d%d", &n, &m, &k, &q);
  int u, v, w;
  foru(i, 1, m) {
    scanf("%d%d%d", &u, &v, &w);
    adj[u].pb(v);
    wei[u].pb(w);
    adj[v].pb(u);
    wei[v].pb(w);
  }
  foru(i, 1, k) {
    scanf("%d", &u);
    fes[u] = 1;
  }
}
/// min path
priority_queue<pii, vector<pii>, greater<pii>> heap;
void dijkstra() {
  foru(u, 1, n) if (fes[u]) {
    dis[u] = 0;
    heap.push(mp(0, u));
  }
  else dis[u] = oo;

  while (!heap.empty()) {
    int u = heap.top().se, e = heap.top().fi;
    heap.pop();
    if (dis[u] != e)
      continue;
    fes[u] = 1;
    foru(i, 0, sz(adj[u]) - 1) {
      int v = adj[u][i];
      if (fes[v])
        continue;
      if (dis[v] > dis[u] + wei[u][i]) {
        dis[v] = dis[u] + wei[u][i];
        heap.push(mp(dis[v], v));
      }
    }
  }
}
/// disjoint set
int root(int u) {
  if (lab[u] == u)
    return u;
  lab[u] = root(lab[u]);
  return lab[u];
}
void disjoint_set() {
  foru(u, 1, n) lab[u] = u;
  int sedge = 0;
  ford(i, m, 1) {
    int u = e[i].se.fi, v = e[i].se.se;
    u = root(u), v = root(v);
    if (u == v)
      continue;
    lab[u] = v;
    sedge++;
    adj[u].pb(v);
    wei[u].pb(e[i].fi);
    adj[v].pb(u);
    wei[v].pb(e[i].fi);
    if (sedge == n - 1)
      break;
  }
}
/// lca
void dfs(int u, int par) {
  foru(i, 0, sz(adj[u]) - 1) {
    int v = adj[u][i];
    if (par == v)
      continue;
    l[v] = l[u] + 1;
    fm[v][0] = wei[u][i];
    p[v][0] = u;
    dfs(v, u);
  }
}
int lca(int u, int v) {
  if (l[u] > l[v])
    swap(u, v);
  if (u == v)
    return 0;
  int res = oo;
  ford(j, 20, 0) if (p[v][j] > 0 && l[p[v][j]] >= l[u]) {
    res = min(res, fm[v][j]);
    v = p[v][j];
  }
  if (u == v)
    return res;
  ford(j, 20, 0) if (p[u][j] > 0 && p[v][j] > 0) if (p[u][j] != p[v][j]) {
    res = min(res, fm[u][j]);
    u = p[u][j];
    res = min(res, fm[v][j]);
    v = p[v][j];
  }
  res = min(res, min(fm[u][0], fm[v][0]));
  return res;
}
/// sol query

int main() {
  // freopen("a.txt", "r", stdin);
  read();
  dijkstra();
  m = 0;
  foru(u, 1, n) {
    foru(i, 0, sz(adj[u]) - 1) {
      int v = adj[u][i];
      e[++m] = mp(min(dis[u], dis[v]), mp(u, v));
    }
    adj[u].clear();
    wei[u].clear();
  }
  sort(e + 1, e + m + 1);
  disjoint_set();

  p[1][0] = -1;
  dfs(1, -1);
  foru(j, 1, 20) foru(u, 1, n) if (p[u][j - 1] > 0) {
    fm[u][j] = min(fm[u][j - 1], fm[p[u][j - 1]][j - 1]);
    p[u][j] = p[p[u][j - 1]][j - 1];
  }
  foru(i, 1, q) {
    int u, v;
    scanf("%d%d", &u, &v);
    printf("%d\\n", lca(u, v));
  }
  return 0;
}