dsa/graphs/trees/ | desc: Least Common Ancestor LCA DP

lca.cpp

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

const int N = 2e5 + 5; // Max num of vertices
const int maxHeight = 20 + 1; // ceil(log2(N)) = 20 for N = 5e5+5
int depth[N], parent[N], dp[N][maxHeight];
vector<int> gph[N];
int n; // Num of vertices (1-indexing)

void dfs (int u, int par, int dep) { // Calculate parent and depth
  parent[u] = par;
  depth[u] = dep;
  for (auto v : gph[u])
    if (v != par) dfs(v, u, dep + 1);
}

void fill_dp(){
  for (int i = 1; i <= n; i++) dp[i][0] = parent[i];
  for (int j = 1; j < maxHeight; j++) {
    for (int i = 1; i <= n; i++) {
      // 8 (2^3) -> 4 (2^2) -> 2 (2^1) -> 1 (2^0)
      dp[i][j] = dp[dp[i][j - 1]][j - 1];
    }
  }
}

int lca(int u, int v) {
  if (depth[v] < depth[u]) swap(u,v); // Make sure u is at lower depth than v
  int diff = depth[v] - depth[u];
  // Bring u and v at same depth by making `v` traverse the height difference
  while (diff > 0) {
    int _log = log2(diff); // Find the greatest power of 2 such that 2^pow <= diff
    v = dp[v][_log];
    diff -= (1 << _log); // Reduce diff by 2*pow
  }
  while (u != v) {
    int _log = log2(depth[u]);
    while (_log > 0 && dp[u][_log] == dp[v][_log]) _log--;
    u = dp[u][_log];
    v = dp[v][_log];
  }
  return u; // u is LCA now
}

int distance_btw_two(int u, int v) {
  return depth[u] + depth[v] - 2 * depth[lca(u, v)];
}

void solve() {
  cin >> n;
  for(int i = 0; i < n - 1; i++){
    int u, v;
    cin >> u >> v;
    gph[u].push_back(v);
    gph[v].push_back(u);
  }
  dfs(1, 0, 0);
  fill_dp();
}