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

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> P;

#define EACH(i, a) for (auto &i : a)
#define FOR(i, a, b) for (ll i = (a); i < (b); i++)
#define RFOR(i, a, b) for (ll i = (b)-1; i >= (a); i--)
#define REP(i, n) for (ll i = 0; i < (n); i++)
#define RREP(i, n) for (ll i = (n)-1; i >= 0; i--)
#define debug(x) cout << #x << ": " << x << endl
#define pb push_back
#define ALL(a) (a).begin(), (a).end()

const ll linf = 1e18;
const int inf = 1e9;
const double eps = 1e-12;
const double pi = acos(-1);

template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
  EACH(x, vec) is >> x;
  return is;
}
/*
template<class... T>
ostream& operator<<(ostream& os, tuple<T...>& t) {
        for (size_t i = 0; i < tuple_size< tuple<T...> >::value; ++i) {
                if (i) os << " ";
                os << get<0>(t);
        }
        return os;
}
*/
template <typename T> ostream &operator<<(ostream &os, vector<T> &vec) {
  REP(i, vec.size()) {
    if (i)
      os << " ";
    os << vec[i];
  }
  return os;
}
template <typename T> ostream &operator<<(ostream &os, vector<vector<T>> &vec) {
  REP(i, vec.size()) {
    if (i)
      os << endl;
    os << vec[i];
  }
  return os;
}

struct Edge {
  int from, to, cost;
};
struct Node {
  int v, cost;
};
bool operator>(const Node &n1, const Node &n2) { return n1.cost > n2.cost; }
bool operator>(const Edge &e1, const Edge &e2) { return e1.cost > e2.cost; }

vector<int> p, h;
class UnionFind {
private:
  int size;

public:
  UnionFind(int size) : size(size) {
    p = vector<int>(size, 0);
    h = vector<int>(size, 0);
    REP(i, size) p[i] = i;
  }
  int root(int n) {
    if (p[n] == n)
      return n;
    return p[n] = root(p[n]);
  }
  bool isUnited(int a, int b) { return root(a) == root(b); }
  void unite(int a, int b) {
    int ra = root(a), rb = root(b);
    if (h[ra] > h[rb]) {
      p[rb] = ra;
    } else if (h[rb] < h[ra]) {
      p[ra] = rb;
    } else {
      p[rb] = ra;
      ++h[ra];
    }
  }
};

int parent[32][100001] = {0};
int depth[100001] = {0};
int mc[32][100001] = {0};

vector<vector<Edge>> G;

void dfs(int v, int prev, int d, int ec) {
  mc[0][v] = min(mc[0][v], ec);
  parent[0][v] = prev;
  depth[v] = d;
  EACH(e, G[v]) {
    if (e.to != prev) {
      dfs(e.to, v, d + 1, e.cost);
    }
  }
}
void init() {
  REP(k, 30) {
    REP(v, G.size()) {
      if (parent[k][v] < 0) {
        mc[k + 1][v] = inf;
        parent[k + 1][v] = -1;
      } else {
        mc[k + 1][v] = min(mc[k][v], mc[k][parent[k][v]]);
        parent[k + 1][v] = parent[k][parent[k][v]];
      }
    }
  }
}
int lca(int u, int v) {
  if (depth[u] > depth[v])
    swap(u, v);
  int res = inf;
  REP(k, 30) {
    if ((depth[v] - depth[u]) & (1LL << k)) {
      res = min(res, mc[k][v]);
      v = parent[k][v];
    }
  }
  if (u == v) {
    //		cout << u+1 << endl;
    return res;
  }
  RREP(k, 30) {
    if (parent[k][u] != parent[k][v]) {
      res = min(res, mc[k][u]);
      res = min(res, mc[k][v]);
      u = parent[k][u];
      v = parent[k][v];
    }
  }
  res = min(res, mc[0][u]);
  res = min(res, mc[0][v]);
  assert(parent[0][u] == parent[0][v]);
  return res; // parent[0][u]
}

signed main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(0);
  int N, M, K, Q;
  cin >> N >> M >> K >> Q;
  vector<Edge> edges;
  G.resize(N);
  REP(i, M) {
    int u, v, d;
    cin >> u >> v >> d;
    --u, --v;
    G[u].pb({u, v, d});
    G[v].pb({v, u, d});
    edges.pb({u, v, d});
  }
  vector<int> dist(N, inf);
  priority_queue<Node, vector<Node>, greater<Node>> pq;
  REP(i, K) {
    int f;
    cin >> f;
    --f;
    pq.push({f, 0});
    dist[f] = 0;
  }
  while (!pq.empty()) {
    Node node = pq.top();
    pq.pop();
    int v = node.v, cost = node.cost;
    if (cost > dist[v])
      continue;
    EACH(e, G[v]) {
      if (cost + e.cost < dist[e.to]) {
        dist[e.to] = cost + e.cost;
        pq.push({e.to, dist[e.to]});
      }
    }
  }
  REP(i, M) {
    Edge &e = edges[i];
    e.cost = min(dist[e.from], dist[e.to]);
  }
  sort(ALL(edges), greater<Edge>());
  UnionFind uf(N);
  G = vector<vector<Edge>>(N);
  int cnt = 0;
  EACH(e, edges) {
    if (!uf.isUnited(e.from, e.to)) {
      uf.unite(e.from, e.to);
      G[e.from].pb({e.from, e.to, e.cost});
      G[e.to].pb({e.to, e.from, e.cost});
      ++cnt;
    }
  }
  assert(cnt == N - 1);
  fill(mc[0], mc[32], inf);
  dfs(0, -1, 0, inf);
  init();

  REP(i, Q) {
    int u, v;
    cin >> u >> v;
    --u, --v;
    cout << lca(u, v) << endl;
  }
}

Error: MLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> P;

#define EACH(i, a) for (auto &i : a)
#define FOR(i, a, b) for (ll i = (a); i < (b); i++)
#define RFOR(i, a, b) for (ll i = (b)-1; i >= (a); i--)
#define REP(i, n) for (ll i = 0; i < (n); i++)
#define RREP(i, n) for (ll i = (n)-1; i >= 0; i--)
#define debug(x) cout << #x << ": " << x << endl
#define pb push_back
#define ALL(a) (a).begin(), (a).end()

const ll linf = 1e18;
const int inf = 1e9;
const double eps = 1e-12;
const double pi = acos(-1);

template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
  EACH(x, vec) is >> x;
  return is;
}
/*
template<class... T>
ostream& operator<<(ostream& os, tuple<T...>& t) {
        for (size_t i = 0; i < tuple_size< tuple<T...> >::value; ++i) {
                if (i) os << " ";
                os << get<0>(t);
        }
        return os;
}
*/
template <typename T> ostream &operator<<(ostream &os, vector<T> &vec) {
  REP(i, vec.size()) {
    if (i)
      os << " ";
    os << vec[i];
  }
  return os;
}
template <typename T> ostream &operator<<(ostream &os, vector<vector<T>> &vec) {
  REP(i, vec.size()) {
    if (i)
      os << endl;
    os << vec[i];
  }
  return os;
}

#define int long long

struct Edge {
  int from, to, cost;
};
struct Node {
  int v, cost;
};
bool operator>(const Node &n1, const Node &n2) { return n1.cost > n2.cost; }
bool operator>(const Edge &e1, const Edge &e2) { return e1.cost > e2.cost; }

vector<int> p, h;
class UnionFind {
private:
  int size;

public:
  UnionFind(int size) : size(size) {
    p = vector<int>(size, 0);
    h = vector<int>(size, 0);
    REP(i, size) p[i] = i;
  }
  int root(int n) {
    if (p[n] == n)
      return n;
    return p[n] = root(p[n]);
  }
  bool isUnited(int a, int b) { return root(a) == root(b); }
  void unite(int a, int b) {
    int ra = root(a), rb = root(b);
    if (h[ra] > h[rb]) {
      p[rb] = ra;
    } else if (h[rb] < h[ra]) {
      p[ra] = rb;
    } else {
      p[rb] = ra;
      ++h[ra];
    }
  }
};

int parent[32][100001] = {0};
int depth[100001] = {0};
int mc[32][100001] = {0};

vector<vector<Edge>> G;

void dfs(int v, int prev, int d, int ec) {
  mc[0][v] = min(mc[0][v], ec);
  parent[0][v] = prev;
  depth[v] = d;
  EACH(e, G[v]) {
    if (e.to != prev) {
      dfs(e.to, v, d + 1, e.cost);
    }
  }
}
void init() {
  REP(k, 30) {
    REP(v, G.size()) {
      if (parent[k][v] < 0) {
        mc[k + 1][v] = linf;
        parent[k + 1][v] = -1;
      } else {
        mc[k + 1][v] = min(mc[k][v], mc[k][parent[k][v]]);
        parent[k + 1][v] = parent[k][parent[k][v]];
      }
    }
  }
}
int lca(int u, int v) {
  if (depth[u] > depth[v])
    swap(u, v);
  int res = linf;
  REP(k, 30) {
    if ((depth[v] - depth[u]) & (1 << k)) {
      res = min(res, mc[k][v]);
      v = parent[k][v];
      assert(v >= 0);
      //			assert( mc[k][v] != linf );
    }
  }
  if (u == v) {
    //		cout << u+1 << endl;
    return res;
  }
  RREP(k, 30) {
    if (parent[k][u] != parent[k][v]) {
      res = min(res, mc[k][u]);
      res = min(res, mc[k][v]);
      u = parent[k][u];
      v = parent[k][v];
    }
  }
  res = min(res, mc[0][u]);
  assert(res != linf);
  //	cout << parent[0][u]+1 << endl;
  return res; // parent[0][u]
}

signed main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(0);
  int N, M, K, Q;
  cin >> N >> M >> K >> Q;
  vector<Edge> edges;
  G.resize(N);
  REP(i, M) {
    int u, v, d;
    cin >> u >> v >> d;
    --u, --v;
    G[u].pb({u, v, d});
    G[v].pb({v, u, d});
    edges.pb({u, v, d});
  }
  vector<int> dist(N, inf);
  priority_queue<Node, vector<Node>, greater<Node>> pq;
  REP(i, K) {
    int f;
    cin >> f;
    --f;
    pq.push({f, 0});
    dist[f] = 0;
  }
  while (!pq.empty()) {
    Node node = pq.top();
    pq.pop();
    int v = node.v, cost = node.cost;
    if (cost > dist[v])
      continue;
    EACH(e, G[v]) {
      if (cost + e.cost < dist[e.to]) {
        dist[e.to] = cost + e.cost;
        pq.push({e.to, dist[e.to]});
      }
    }
  }
  REP(i, M) {
    Edge &e = edges[i];
    e.cost = min(dist[e.from], dist[e.to]);
  }
  sort(ALL(edges), greater<Edge>());
  UnionFind uf(N);
  G = vector<vector<Edge>>(N);
  int cnt = 0;
  EACH(e, edges) {
    if (!uf.isUnited(e.from, e.to)) {
      uf.unite(e.from, e.to);
      G[e.from].pb({e.from, e.to, e.cost});
      G[e.to].pb({e.to, e.from, e.cost});
      ++cnt;
    }
  }
  assert(cnt == N - 1);
  fill(mc[0], mc[32], linf);
  dfs(0, -1, 0, linf);
  init();

  REP(i, Q) {
    int u, v;
    cin >> u >> v;
    --u, --v;
    cout << lca(u, v) << endl;
  }
}