Directory

• Preprocessing
• LCA
• External Function Version

Preprocessing

void solve() {
    int n, k;
    cin >> n >> k;
    vector<vector<int>> adj(n + 1);
    for (int i = 1; i <= n - 1; ++i) {
        int x, y;
        cin >> x >> y;
        adj[x].push_back(y);
        adj[y].push_back(x);
    }
    vector<vector<int>> parent(n + 1, vector<int>(22));
    vector<int> depth(n + 1);
    vector<int> leaves;

    auto dfs = [&](int cur, int pa, auto&& dfs) -> void {
        depth[cur] = depth[pa] + 1;
        if (adj[cur].size() == 1) {
            leaves.push_back(cur);
        }
        for (auto& nxt : adj[cur]) {
            if (nxt != pa) {
                parent[nxt][0] = cur;
                dfs(nxt, cur, dfs);
            }
        }
    };
    dfs(1, 0, dfs);

    // Binary lifting to compute ancestors
    for (int p = 1; p < 22; ++p) {
        for (int i = 1; i <= n; ++i) {
            parent[i][p] = parent[parent[i][p - 1]][p - 1];
        }
    }
}

LCA

auto LCA = [&](int a, int b) -> pair<int, int> {
    long long ans = 0;
    if (depth[a] < depth[b]) swap(a, b);

    // Lift a to depth of b
    while (depth[a] > depth[b]) {
        int diff = depth[a] - depth[b];
        int step = (int)log2(diff);
        a = parent[a][step];
        ans += (1LL << step);
    }

    // Lift both together until they meet
    for (int i = log2(depth[a]); i >= 0; --i) {
        if (parent[a][i] != parent[b][i]) {
            a = parent[a][i];
            b = parent[b][i];
            ans += (1LL << (i + 1));
        }
    }
    if (a != b) {
        a = parent[a][0];
        b = parent[b][0];
        ans += 2;
    }
    return {a, ans};
};

The first while loop lifts a to the same depth as b using binary decomposition.

Etxernal Function Version

inline int LCA(int a, int b) {
    if (depth[a] < depth[b]) swap(a, b);
    while (depth[a] > depth[b]) {
        int step = (int)log2(depth[a] - depth[b]);
        a = parent[a][step];
    }
    if (a == b) return a;
    for (int i = log2(depth[a]); i >= 0; --i) {
        if (parent[a][i] != parent[b][i]) {
            a = parent[a][i];
            b = parent[b][i];
        }
    }
    if (a != b) {
        a = parent[a][0];
        b = parent[b][0];
    }
    return a;
}

Usage Example with Forward Star

(For a complete example with forward star storage, see the reference: POJ 3417 Network – LCA with tree difference and forward star needing fast input)

struct Edge {
    int to;
    int next;
} edges[2 * maxn];
int head[maxn];
int total = 0;

inline void addEdge(int from, int to) {
    ++total;
    edges[total].to = to;
    edges[total].next = head[from];
    head[from] = total;
}

int parent[maxn][23];
int depth[maxn];

void dfs1(int cur, int pa) {
    depth[cur] = depth[pa] + 1;
    parent[cur][0] = pa;
    for (int p = 1; p < 22; ++p)
        parent[cur][p] = parent[parent[cur][p - 1]][p - 1];

    for (int i = head[cur]; i > 0; i = edges[i].next) {
        int nxt = edges[i].to;
        if (nxt != pa)
            dfs1(nxt, cur);
    }
}

inline int LCA(int a, int b) {
    if (depth[a] < depth[b]) swap(a, b);
    while (depth[a] > depth[b]) {
        int step = (int)log2(depth[a] - depth[b]);
        a = parent[a][step];
    }
    if (a == b) return a;
    for (int i = log2(depth[a]); i >= 0; --i) {
        if (parent[a][i] != parent[b][i]) {
            a = parent[a][i];
            b = parent[b][i];
        }
    }
    if (a != b) {
        a = parent[a][0];
        b = parent[b][0];
    }
    return a;
}

int diff[maxn];
int pass[maxn];
int cnt0, cnt1;

void dfs2(int cur, int pa) {
    pass[cur] += diff[cur];
    for (int i = head[cur]; i > 0; i = edges[i].next) {
        int nxt = edges[i].to;
        if (nxt != pa) {
            dfs2(nxt, cur);
            pass[cur] += pass[nxt];
        }
    }
}