Company Queries II

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

const int MAXN = 300000;
const int LOG = 18;

vector<int> TREE[MAXN];
int LCA[MAXN][LOG];
int DEPTH[MAXN];

// Find all node depth and fill LCA table
void dfs(int a, int b) {
    LCA[a][0] = b;

    // Fill LCA
    for (int i = 1; i < LOG; i++) {
        if (LCA[a][i - 1] != -1) {
            LCA[a][i] = LCA[LCA[a][i - 1]][i - 1];
        } else {
            LCA[a][i] = -1;
        }
    }

    // Go thru children and set depth
    for (int u : TREE[a]) {
        if (u != b) {
            DEPTH[u] = DEPTH[a] + 1;
            dfs(u, a);
        }
    }
}

// Find LCA of a and b
int lca(int a, int b) {
    if (DEPTH[a] < DEPTH[b]) {
        swap(a, b);
    }

    // Get depth difference and fix it
    int diff = DEPTH[a] - DEPTH[b];
    for (int i = 0; i < LOG; i++) {
        if ((diff >> i) & 1) {
            a = LCA[a][i];
        }
    }

    // If the node was on the same branch, check if it is now b
    if (a == b) {
        return a;
    }

    // Both are on the same level, keep moving them up, until they are the same.
    for (int i = LOG - 1; i >= 0; i--) {
        if (LCA[a][i] != LCA[b][i]) {
            a = LCA[a][i];
            b = LCA[b][i];
        }
    }

    return LCA[a][0];
}

int main() {
    int n, q;
    cin >> n >> q; // input(), get from standard input feed

    // Build tree
    for (int child = 2; child <= n; child++) {
        int boss;
        cin >> boss;
        TREE[boss].push_back(child);
        TREE[child].push_back(boss);
    }

    // Queries
    vector<vector<int> > queries(q, vector<int>(2));
    for (int query = 0; query < q; query++) {
        int a, b;
        cin >> a >> b;

        queries[query][0] = a;
        queries[query][1] = b;
    }

    DEPTH[1] = 0;
    
    // Run DFS
    dfs(1, -1);

    for (int i = 0; i < q; ++i) {
        int a = queries[i][0];
        int b = queries[i][1];
        cout << lca(a, b) << "\n";
    }

    return 0;
}

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: ACCEPTED

Test 4

Verdict: ACCEPTED

Test 5

Verdict: ACCEPTED

Test 6

Verdict: ACCEPTED

Test 7

Verdict: ACCEPTED

Test 8

Verdict: ACCEPTED

Test 9

Verdict: ACCEPTED

Test 10

Verdict: ACCEPTED

Test 11

Verdict: ACCEPTED

Test 12

Verdict: ACCEPTED