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

#include <unordered_map>
#include <utility>
template <typename T> class UnionFind {
public:
  void unite(const T &a, const T &b) {
    T x = find(a), y = find(b);
    if (rank_[x] < rank_[y])
      std::swap(x, y);
    parent_[y] = x;
    if (rank_[x] == rank_[y])
      rank_[x] += 1;
  }
  const T find(const T &x) {
    if (!parent_.count(x)) {
      rank_[x] = 1;
      return parent_[x] = x;
    }
    return (parent_[x] == x) ? x : (parent_[x] = find(parent_[x]));
  }

private:
  std::unordered_map<T, T> parent_;
  std::unordered_map<T, int> rank_;
};

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

struct Query {
  int S, T;
};

struct Edge {
  int from, to, weight;
};

typedef vector<vector<Edge>> Graph;

const vector<int> Dijkstra(const Graph &G, const vector<int> &start) {
  vector<int> d(G.size(), -1);
  priority_queue<tuple<int, int>, vector<tuple<int, int>>,
                 greater<tuple<int, int>>>
      q;
  for (auto v : start)
    q.emplace(0, v);
  while (!q.empty()) {
    int cost, v;
    tie(cost, v) = q.top();
    q.pop();
    if (~d[v])
      continue;
    d[v] = cost;
    for (const auto &e : G[v])
      q.emplace(cost + e.weight, e.to);
  }
  return d;
}

// Kruskal
const vector<int> solve(const Graph &G, const vector<int> &F,
                        const vector<Query> &query) {
  vector<int> res(query.size());
  vector<set<int>> qset(G.size());
  for (int i = 0; i < query.size(); ++i) {
    qset[query[i].S].insert(i);
    qset[query[i].T].insert(i);
  }
  auto d = Dijkstra(G, F);
  set<Edge, function<bool(Edge, Edge)>> E([](Edge lhs, Edge rhs) {
    if (lhs.weight != rhs.weight)
      return lhs.weight > rhs.weight;
    if (lhs.from != rhs.from)
      return lhs.from < rhs.from;
    return lhs.to < rhs.to;
  });
  for (const auto &list : G)
    for (const auto &e : list)
      E.insert({e.from, e.to, min(d[e.from], d[e.to])});
  UnionFind<int> group;
  while (!E.empty()) {
    auto e = *begin(E);
    E.erase(begin(E));
    int u = group.find(e.from);
    int v = group.find(e.to);
    if (u == v)
      continue;
    if (qset[u].size() > qset[v].size())
      swap(u, v);
    for (const auto &q : qset[u]) {
      if (qset[v].count(q) && !res[q])
        res[q] = e.weight;
      qset[v].insert(q);
    }
    group.unite(u, v);
    if (group.find(u) == v) {
      qset[u].clear();
    } else {
      qset[u] = qset[v];
      qset[v].clear();
    }
  }
  return res;
}

int main() {
  int N, M, K, Q;
  cin >> N >> M >> K >> Q;

  Graph G(N + 1);
  for (int i = 0; i < M; ++i) {
    int A, B, L;
    cin >> A >> B >> L;
    G[A].push_back({A, B, L});
    G[B].push_back({B, A, L});
  }

  vector<int> F(K);
  for (auto &i : F)
    cin >> i;

  vector<Query> query(Q);
  for (auto &q : query)
    cin >> q.S >> q.T;

  for (const auto &i : solve(G, F, query))
    cout << i << endl;
}

Error: MLE

Fixed submission:

#include <utility>
#include <vector>
class UnionFind {
public:
  UnionFind(int n) : parent_(n, -1), rank_(n) {}
  void unite(int a, int b) {
    int x = find(a), y = find(b);
    if (rank_[x] < rank_[y])
      std::swap(x, y);
    parent_[y] = x;
    if (rank_[x] == rank_[y])
      rank_[x] += 1;
  }
  int find(int x) {
    if (~parent_[x])
      return (parent_[x] == x) ? x : (parent_[x] = find(parent_[x]));
    rank_[x] = 1;
    return parent_[x] = x;
  }

private:
  std::vector<int> parent_;
  std::vector<int> rank_;
};

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

struct Query {
  int S, T;
};

struct Edge {
  int from, to, weight;
};

typedef vector<vector<Edge>> Graph;

const vector<int> Dijkstra(const Graph &G, const vector<int> &start) {
  vector<int> d(G.size(), -1);
  priority_queue<tuple<int, int>, vector<tuple<int, int>>,
                 greater<tuple<int, int>>>
      q;
  for (auto v : start)
    q.emplace(0, v);
  while (!q.empty()) {
    int cost, v;
    tie(cost, v) = q.top();
    q.pop();
    if (~d[v])
      continue;
    d[v] = cost;
    for (const auto &e : G[v])
      q.emplace(cost + e.weight, e.to);
  }
  return d;
}

// Kruskal
const vector<int> solve(const Graph &G, const vector<int> &F,
                        const vector<Query> &query) {
  vector<int> res(query.size());
  vector<set<int>> qset(G.size());
  for (int i = 0; i < query.size(); ++i) {
    qset[query[i].S].insert(i);
    qset[query[i].T].insert(i);
  }
  auto d = Dijkstra(G, F);
  set<Edge, function<bool(Edge, Edge)>> E([](Edge lhs, Edge rhs) {
    if (lhs.weight != rhs.weight)
      return lhs.weight > rhs.weight;
    if (lhs.from != rhs.from)
      return lhs.from < rhs.from;
    return lhs.to < rhs.to;
  });
  for (const auto &list : G)
    for (const auto &e : list)
      E.insert({e.from, e.to, min(d[e.from], d[e.to])});
  UnionFind group(G.size() + 1);
  while (!E.empty()) {
    auto e = *begin(E);
    E.erase(begin(E));
    int u = group.find(e.from);
    int v = group.find(e.to);
    if (u == v)
      continue;
    if (qset[u].size() > qset[v].size())
      swap(u, v);
    for (const auto &q : qset[u]) {
      if (qset[v].count(q) && !res[q])
        res[q] = e.weight;
      qset[v].insert(q);
    }
    group.unite(u, v);
    qset[group.find(u)] = qset[v];
  }
  return res;
}

int main() {
  int N, M, K, Q;
  cin >> N >> M >> K >> Q;

  Graph G(N + 1);
  for (int i = 0; i < M; ++i) {
    int A, B, L;
    cin >> A >> B >> L;
    G[A].push_back({A, B, L});
    G[B].push_back({B, A, L});
  }

  vector<int> F(K);
  for (auto &i : F)
    cin >> i;

  vector<Query> query(Q);
  for (auto &q : query)
    cin >> q.S >> q.T;

  for (const auto &i : solve(G, F, query))
    cout << i << endl;
}