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

#include <bits/stdc++.h>

using namespace std;

#define int long long

#define rep(i, x) for (int i = 0; i < x; ++i)
#define rrep(i, x) for (int i = x - 1; i >= 0; --i)
#define rep1(i, x) for (int i = 1; i <= x; ++i)
#define PB push_back
#define PPB pop_back
#define fst first
#define scd second

using vi = vector<int>;
using vvi = vector<vi>;
using pii = pair<int, int>;
using vpii = vector<pii>;

struct edge {
  int to, cost;
};
using Graph = vector<edge>;

const int inf = 3e18;

struct UF {
  vi data;
  UF(int n) { data.resize(n, -1); }
  int root(int i) { return data[i] < 0 ? i : data[i] = root(data[i]); }
  int size(int x) {
    x = root(x);
    return -data[x];
  }
  void unite(int x, int y) {
    x = root(x), y = root(y);
    if (x == y)
      return;
    if (size(x) < size(y))
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
  }
  bool same(int x, int y) {
    x = root(x), y = root(y);
    return x == y;
  }
};

Graph G[100010];
int d[100010];
int N, M, K, Q;
int F[100010];

// LCA??¨
constexpr int max_log_n = log2(100010);
vi tr[100010];
int parent[100010][max_log_n + 5];
int minimum[100010][max_log_n + 5];
int depth[100010];

void lca_dfs(int cur = 0, int par = -1, int dep = 0) {
  parent[cur][0] = par;
  minimum[cur][0] = min(d[cur], par == -1 ? inf : d[par]);
  depth[cur] = dep;

  for (int nx : tr[cur]) {
    if (nx == par)
      continue;
    lca_dfs(nx, cur, dep + 1);
  }
}

void lca_init() {
  lca_dfs();

  rep(i, max_log_n) {
    rep(j, N) {
      if (~parent[j][i]) {
        parent[j][i + 1] = parent[parent[j][i]][i];
        minimum[j][i + 1] = min(minimum[j][i], minimum[parent[j][i]][i]);
      } else {
        parent[j][i + 1] = -1;
        minimum[j][i + 1] = inf;
      }
    }
  }
}

int lca(int u, int v) {
  if (depth[u] > depth[v])
    swap(u, v);

  int diff = depth[v] - depth[u];

  rep(i, max_log_n + 1) {
    if (diff >> i & 1)
      v = parent[v][i];
  }

  if (u == v)
    return u;

  rrep(i, max_log_n + 1) {
    if (parent[u][i] != parent[v][i]) {
      u = parent[u][i];
      v = parent[v][i];
    }
  }

  return parent[u][0];
}

int calc2(int p, int c) {
  if (p == c)
    return d[p];

  int dif = depth[c] - depth[p];

  int mi = d[c];

  rep(i, max_log_n + 1) {
    if (dif & 1 << i) {
      mi = min(mi, minimum[c][i]);
      c = parent[c][i];
    }
  }

  return mi;
}

int calc(int u, int v) {
  if (u == v)
    return d[u];

  int l = lca(u, v);

  u = calc2(l, u), v = calc2(l, v);

  return min(u, v);
}

void dijkstra() {
  priority_queue<pii, vpii, greater<pii>> q;
  rep(i, K) {
    d[F[i]] = 0;
    q.push(pii(0, F[i]));
  }
  while (q.size()) {
    pii p = q.top();
    q.pop();
    int v = p.scd;
    if (d[v] < p.fst)
      continue;
    for (edge &e : G[v]) {
      if (d[e.to] > d[v] + e.cost) {
        d[e.to] = d[v] + e.cost;
        q.push(pii(d[e.to], e.to));
      }
    }
  }
}

void build_tree() {
  using EdgeInfo = pair<int, pair<int, int>>;
  vector<EdgeInfo> E;

  rep(i, N) {
    for (edge &e : G[i]) {
      E.PB(EdgeInfo{min(d[i], d[e.to]), pii(i, e.to)});
    }
  }

  sort(begin(E), end(E), greater<EdgeInfo>());

  UF uf(N);

  int n = E.size();
  rep(i, n) {
    int from = E[i].scd.fst, to = E[i].scd.scd;

    if (uf.same(from, to))
      continue;
    uf.unite(from, to);

    tr[from].PB(to);
    tr[to].PB(from);
  }
}

signed main() {
  cin >> N >> M >> K >> Q;

  rep(i, M) {
    int a, b, l;
    cin >> a >> b >> l;
    --a, --b;

    G[a].PB(edge{b, l});
    G[b].PB(edge{a, l});
  }

  rep(i, K) {
    cin >> F[i];
    --F[i];
  }

  fill_n(d, N + 1, inf);
  dijkstra();

  build_tree();

  lca_init();

  rep(i, Q) {
    int s, t;
    cin >> s >> t;
    --s, --t;
    int v = calc(s, t);
    cout << v << endl;
  }

  // rep(i, N) {
  //     cout << "d[" << i+1 << "] = " << d[i] << endl;
  // }

  // cout << "tree_debug:\\n";
  // rep(v, N) {
  //     for (int u : tr[v]) {
  //         if (v < u) {
  //             cout << v+1 << " <-> " << u+1 << endl;
  //         }
  //     }
  // }
}

Error: MLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

#define int long long

#define rep(i, x) for (int i = 0; i < x; ++i)
#define rrep(i, x) for (int i = x - 1; i >= 0; --i)
#define rep1(i, x) for (int i = 1; i <= x; ++i)
#define PB push_back
#define PPB pop_back
#define fst first
#define scd second

using vi = vector<int>;
using vvi = vector<vi>;
using pii = pair<int, int>;
using vpii = vector<pii>;

struct edge {
  int to, cost;
};
using Graph = vector<edge>;

const int inf = 1e9;

struct UF {
  vi data;
  UF(int n) { data.resize(n, -1); }
  int root(int i) { return data[i] < 0 ? i : data[i] = root(data[i]); }
  int size(int x) {
    x = root(x);
    return -data[x];
  }
  void unite(int x, int y) {
    x = root(x), y = root(y);
    if (x == y)
      return;
    if (size(x) < size(y))
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
  }
  bool same(int x, int y) {
    x = root(x), y = root(y);
    return x == y;
  }
};

Graph G[100010];
int d[100010];
int N, M, K, Q;
int F[100010];

// LCA??¨
vi tr[100010];
int depth[100010];
int parent[100010]
          [30]; // parent[i][j] := ??????i???2^j??????????????????????????????,
                // ?????¨???????????´??????-1
int minimum[100010][30];

void lca_dfs(int cur = 0, int p = -1, int dep = 0) {
  depth[cur] = dep;
  parent[cur][0] = p;
  minimum[cur][0] = min(d[cur], p == -1 ? inf : d[p]);

  for (int &u : tr[cur]) {
    if (u == p)
      continue;
    lca_dfs(u, cur, dep + 1);
  }
}

void lca_init() {
  lca_dfs();

  vpii hoge;

  rep(i, N) { hoge.PB(pii(depth[i], i)); }

  sort(begin(hoge), end(hoge));

  rep(x, 20) {
    rep(ii, N) {
      int i = hoge[ii].scd;
      if (~parent[i][x]) {
        parent[i][x + 1] = parent[parent[i][x]][x];
        minimum[i][x + 1] = min(minimum[i][x], minimum[parent[i][x]][x]);
      } else {
        parent[i][x + 1] = -1;
        minimum[i][x + 1] = inf;
      }
    }
  }
}

int lca(int u, int v) {
  if (depth[u] < depth[v])
    swap(u, v);

  int dif = depth[u] - depth[v];

  rep(i, 20) {
    if (dif & 1 << i) {
      u = parent[u][i];
    }
  }

  if (u == v)
    return u;

  rrep(i, 21) {
    if (parent[u][i] != parent[v][i]) {
      u = parent[u][i];
      v = parent[v][i];
    }
  }

  return parent[u][0];
}

int calc2(int p, int c) {
  if (p == c)
    return d[p];

  int dif = depth[c] - depth[p];

  int mi = d[c];

  rep(i, 20) {
    if (dif & 1 << i) {
      mi = min(mi, minimum[c][i]);
      c = parent[c][i];
    }
  }

  return mi;
}

int calc(int u, int v) {
  if (u == v)
    return d[u];

  int l = lca(u, v);

  int a = calc2(l, u), b = calc2(l, v);

  return min(a, b);
}

void dijkstra() {
  priority_queue<pii, vpii, greater<pii>> q;
  rep(i, K) {
    d[F[i]] = 0;
    q.push(pii(0, F[i]));
  }
  while (q.size()) {
    pii p = q.top();
    q.pop();
    int v = p.scd;
    if (d[v] < p.fst)
      continue;
    for (edge &e : G[v]) {
      if (d[e.to] > d[v] + e.cost) {
        d[e.to] = d[v] + e.cost;
        q.push(pii(d[e.to], e.to));
      }
    }
  }
}

void build_tree() {
  using EdgeInfo = pair<pair<int, int>, pair<int, int>>;
  vector<EdgeInfo> E;

  rep(i, N) {
    for (edge &e : G[i]) {
      E.PB(EdgeInfo{pii(d[i], d[e.to]), pii(i, e.to)});
    }
  }

  sort(begin(E), end(E), greater<EdgeInfo>());

  UF uf(N);

  int n = E.size();
  rep(i, n) {
    int from = E[i].scd.fst, to = E[i].scd.scd;

    if (uf.same(from, to))
      continue;
    uf.unite(from, to);

    tr[from].PB(to);
    tr[to].PB(from);
  }
}

signed main() {
  cin >> N >> M >> K >> Q;

  rep(i, M) {
    int a, b, l;
    cin >> a >> b >> l;
    --a, --b;

    G[a].PB(edge{b, l});
    G[b].PB(edge{a, l});
  }

  rep(i, K) {
    cin >> F[i];
    --F[i];
  }

  fill_n(d, N + 1, inf);
  dijkstra();

  build_tree();

  lca_init();

  rep(i, Q) {
    int s, t;
    cin >> s >> t;
    --s, --t;
    int v = calc(s, t);
    cout << v << endl;
  }

  // rep(i, N) {
  //     cout << "d[" << i+1 << "] = " << d[i] << endl;
  // }

  // cout << "tree_debug:\\n";
  // rep(v, N) {
  //     for (int u : tr[v]) {
  //         if (v < u) {
  //             cout << v+1 << " <-> " << u+1 << endl;
  //         }
  //     }
  // }
}