Company Queries II

#include <iostream>
#include <vector>
#include <utility>

int ancestor(int node, int depth, std::vector<std::vector<int>> &boss_of_powers) {
    int power = 0;
    int number = 1;
    while (number <= depth) {
        power++;
        number *= 2;
    }
    
    for (int i = power-1; i >= 0; --i) {
        number /= 2;
        if (depth >= number) {
            node = boss_of_powers[node][i];
            depth -= number;
        }
    }
    return node;
}

int find_common_ancestor(int node1, int node2, int max_depth, std::vector<std::vector<int>> &boss_of_powers) {
    if (node1 == node2) return node1;
    if (boss_of_powers[node1][0] == boss_of_powers[node2][0]) return boss_of_powers[node1][0];

    int power = 0;
    int number = 1;
    while (number <= max_depth) {
        power++;
        number *= 2;
    }

    for (int i = power-1; i >= 0; --i) {
        number /= 2;
        if (boss_of_powers[node1][i] != boss_of_powers[node2][i]) {
            return find_common_ancestor(boss_of_powers[node1][i], boss_of_powers[node2][i], max_depth-number, boss_of_powers);
        }
    }
    return 0;
}

int main() {
    int n, q;
    std::cin >> n >> q;

    std::vector<std::vector<int>> boss_of_powers(n+1, std::vector(1, 0));
    for (int i = 2; i <= n; ++i) {
        std::cin >> boss_of_powers[i][0];
    }
    int power = 1;
    int number = 2;
    while (number <= n) {
        for (int i = 1; i <= n; ++i) {
            int boss = boss_of_powers[boss_of_powers[i][power-1]][power-1];
            boss_of_powers[i].push_back(boss);
        }

        power++;
        number *= 2;
    }

    std::vector<int> depths(n+1, 0);
    for (int i = 2; i <= n; ++i) {
        depths[i] = depths[boss_of_powers[i][0]] + 1;
    }

    std::vector<std::pair<int, int>> queries(q);
    for (auto &query : queries) {
        std::cin >> query.first >> query.second;
    }
    
    std::vector<int> results(q);
    for (int i = 0; i < q; ++i) {
        int node1 = queries[i].first;
        int node2 = queries[i].second;
        
        //std::cout << node1 << " " << node2 << " " << depths[node1] << " " << depths[node2] << std::endl;
        if (depths[node1] < depths[node2]) {
            node2 = ancestor(node2, depths[node2] - depths[node1], boss_of_powers);
        }
        if (depths[node1] > depths[node2]) {
            node1 = ancestor(node1, depths[node1] - depths[node2], boss_of_powers);
        }

        //std::cout << node1 << " " << node2 << " " << depths[node1] << std::endl;

        results[i] = find_common_ancestor(node1, node2, depths[node1], boss_of_powers);
    }

    for (auto &result : results) {
        std::cout << result << std::endl;
    }
}

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: ACCEPTED

Test 4

Verdict: ACCEPTED

Test 5

Verdict: ACCEPTED

Test 6

Verdict: RUNTIME ERROR

Test 7

Verdict: ACCEPTED

Test 8

Verdict: ACCEPTED

Test 9

Verdict: RUNTIME ERROR

Test 10

Verdict: RUNTIME ERROR

Test 11

Verdict: ACCEPTED

Test 12

Verdict: RUNTIME ERROR