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

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

struct UnionFind {
  vector<int> data;
  UnionFind(int sz) { data.resize(sz, -1); }
  int Find(int a) { return (data[a] < 0 ? a : data[a] = Find(data[a])); }
  void Union(int a, int b) {
    a = Find(a), b = Find(b);
    if (a == b)
      return;
    if (data[a] < data[b])
      swap(a, b);
    data[b] += data[a];
    data[a] = b;
  }
};

struct edge {
  int to, cost;
};
struct edge2 {
  int u, v, cost;
  bool operator<(const edge2 &e) const { return cost < e.cost; }
};
typedef vector<vector<edge>> Graph;
typedef pair<int, int> Pi;
const int INF = 1 << 30;
int N, M, K, Q, F[100000];
Graph graph;
vector<edge2> weighted;
set<int> buffer[100000];
stack<int> st;
int ret[100000];

void Dijkstra(Graph &g, vector<int> &min_cost) {
  min_cost.assign(g.size(), INF);
  priority_queue<Pi, vector<Pi>, greater<Pi>> que;
  for (int i = 0; i < K; i++) {
    que.push(make_pair(0, F[i]));
    min_cost[F[i]] = 0;
  }
  while (!que.empty()) {
    Pi p = que.top();
    que.pop();
    if (p.first > min_cost[p.second])
      continue;
    for (int i = 0; i < g[p.second].size(); i++) {
      edge &e = g[p.second][i];
      if (p.first + e.cost >= min_cost[e.to])
        continue;
      que.push(make_pair(p.first + e.cost, e.to));
      min_cost[e.to] = p.first + e.cost;
    }
  }
}

int main() {
  scanf("%d %d %d %d", &N, &M, &K, &Q);
  graph.resize(N);
  UnionFind uftree(N);

  for (int i = 0; i < M; i++) {
    int a, b, c;
    scanf("%d %d %d", &a, &b, &c);
    --a, --b;
    graph[a].push_back((edge){b, c});
    graph[b].push_back((edge){a, c});
    weighted.push_back((edge2){a, b});
  }

  for (int i = 0; i < K; i++) {
    scanf("%d", &F[i]);
    --F[i];
  }
  for (int i = 0; i < Q; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    --a, --b;
    buffer[a].insert(i);
    buffer[b].insert(i);
  }

  vector<int> min_cost;
  Dijkstra(graph, min_cost);

  for (int i = 0; i < M; i++) {
    weighted[i].cost = min(min_cost[weighted[i].u], min_cost[weighted[i].v]);
  }
  sort(weighted.rbegin(), weighted.rend());

  for (int i = 0; i < M; i++) {
    int u = uftree.Find(weighted[i].u), v = uftree.Find(weighted[i].v);
    if (u == v)
      continue;
    uftree.Union(u, v);
    int unioned = uftree.Find(weighted[i].u);
    set<int> *oo = &buffer[u], *pp = &buffer[v];
    if (u != unioned)
      swap(oo, pp);
    for (set<int>::iterator it = pp->begin(); it != pp->end(); it++) {
      if (oo->find(*it) != oo->end())
        ret[*it] = weighted[i].cost;
      else
        st.push(*it);
    }
    for (; !st.empty(); oo->insert(st.top()), st.pop())
      ;
  }

  for (int i = 0; i < Q; i++) {
    printf("%d\\n", ret[i]);
  }
}

Error: MLE

Fixed submission:

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

struct UnionFind {
  vector<int> data;
  UnionFind(int sz) { data.resize(sz, -1); }
  int Find(int a) { return (data[a] < 0 ? a : data[a] = Find(data[a])); }
  void Union(int a, int b) {
    a = Find(a), b = Find(b);
    if (a == b)
      return;
    if (data[a] > data[b])
      swap(a, b);
    data[b] += data[a];
    data[a] = b;
  }
};

struct edge {
  int to, cost;
};
struct edge2 {
  int u, v, cost;
  bool operator<(const edge2 &e) const { return cost < e.cost; }
};
typedef vector<vector<edge>> Graph;
typedef pair<int, int> Pi;
const int INF = 1 << 30;
int N, M, K, Q, F[100000];
Graph graph;
vector<edge2> weighted;
set<int> buffer[100000];
stack<int> st;
int ret[100000];

void Dijkstra(Graph &g, vector<int> &min_cost) {
  min_cost.assign(g.size(), INF);
  priority_queue<Pi, vector<Pi>, greater<Pi>> que;
  for (int i = 0; i < K; i++) {
    que.push(make_pair(0, F[i]));
    min_cost[F[i]] = 0;
  }
  while (!que.empty()) {
    Pi p = que.top();
    que.pop();
    if (p.first > min_cost[p.second])
      continue;
    for (int i = 0; i < g[p.second].size(); i++) {
      edge &e = g[p.second][i];
      if (p.first + e.cost >= min_cost[e.to])
        continue;
      que.push(make_pair(p.first + e.cost, e.to));
      min_cost[e.to] = p.first + e.cost;
    }
  }
}

int main() {
  scanf("%d %d %d %d", &N, &M, &K, &Q);
  graph.resize(N);
  UnionFind uftree(N);

  for (int i = 0; i < M; i++) {
    int a, b, c;
    scanf("%d %d %d", &a, &b, &c);
    --a, --b;
    graph[a].push_back((edge){b, c});
    graph[b].push_back((edge){a, c});
    weighted.push_back((edge2){a, b});
  }

  for (int i = 0; i < K; i++) {
    scanf("%d", &F[i]);
    --F[i];
  }
  for (int i = 0; i < Q; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    --a, --b;
    buffer[a].insert(i);
    buffer[b].insert(i);
  }

  vector<int> min_cost;
  Dijkstra(graph, min_cost);

  for (int i = 0; i < M; i++) {
    weighted[i].cost = min(min_cost[weighted[i].u], min_cost[weighted[i].v]);
  }
  sort(weighted.rbegin(), weighted.rend());

  for (int i = 0; i < M; i++) {
    int u = uftree.Find(weighted[i].u), v = uftree.Find(weighted[i].v);
    if (u == v)
      continue;
    uftree.Union(u, v);
    int unioned = uftree.Find(weighted[i].u);
    set<int> *oo = &buffer[u], *pp = &buffer[v];
    if (u != unioned)
      swap(oo, pp);
    for (set<int>::iterator it = pp->begin(); it != pp->end(); it++) {
      if (oo->find(*it) != oo->end())
        ret[*it] = weighted[i].cost;
      else
        st.push(*it);
    }
    for (; !st.empty(); oo->insert(st.top()), st.pop())
      ;
  }

  for (int i = 0; i < Q; i++) {
    printf("%d\\n", ret[i]);
  }
}