Company Queries II

#include <vector>
#include <iostream>
#include <math.h>
using namespace std;

vector<vector<int>> ancestors;  // a[node][lb(x)]
vector<int> depth;
vector<bool> been;
vector<int> boss;

int dfs(int node) {
    if (been[node]) return depth[node];
    been[node] = true;
    if (node == 0) return depth[node];
    int d = dfs(boss[node]) + 1;
    depth[node] = d;
    return d;
}

int ancestor(int node, int up) {
    for (int ind=0; ind<(int)ancestors[node].size(); ind++) {
        if ((1 << ind) & up) {
            node = ancestors[node][ind];
        }
    }
    return node;
}

int main() {
    int n, q;
    cin >> n >> q;
    ancestors.resize(n);
    depth.resize(n);
    been.resize(n);
    boss.resize(n);
    for (int i=0; i<n; i++) { been[i] = false; }

    boss[0] = 0;
    depth[0] = 0;
    
    int a, b;
    for (int i=1; i<n; i++) {
        cin >> b;
        b--;
        boss[i] = b;
    }

    for (int i=0; i<n; i++) { dfs(i); }

    int logN = log2(2*n);
    
    // set ancestors
    for (int i=0; i<n; i++) { ancestors[i].push_back(boss[i]); }
    for (int i=0; i<logN; i++) {
        for (int j=0; j<n; j++) {
            ancestors[j].push_back(ancestors[ancestors[j][i]][i]);
        }
    }
    int da, db;
    for (int i=0; i<q; i++) {
        cin >> a >> b;
        a--; b--;
        da = depth[a];
        db = depth[b];
        if (da > db) swap(a, b);  // b is lower
        b = ancestor(b, abs(da - db));  // level

        if (a == b) {
            cout << a + 1 << endl;
            continue;
        }

        for (int k=logN-1; k>=0; k--) {
            if (ancestors[a][k] != ancestors[b][k]) {
                a = ancestors[a][k];
                b = ancestors[b][k];
            }
        }
        cout << ancestor(a, 1) + 1 << endl;
    }

    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: TIME LIMIT EXCEEDED