Company Queries II

#![allow(unused)]

use std::io;
fn take_tuple() -> (usize, usize) {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap();
    let mut it = input
        .trim()
        .split_whitespace()
        .map(|x| x.parse::<usize>().unwrap());
    (it.next().unwrap(), it.next().unwrap())
}
fn take_vector() -> Vec<usize> {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap();
    let arr: Vec<usize> = input
        .trim()
        .split_whitespace()
        .map(|x| x.parse::<usize>().unwrap())
        .collect();
    return arr;
}

fn main() {
    let (n, q): (usize, usize) = take_tuple();
    let mut vec = take_vector();
    let max_tree: usize = usize::max(n.next_power_of_two().ilog2() as usize, 1);

    let mut ancestors: Vec<Vec<usize>> = vec![Vec::with_capacity(n + 1); max_tree];
    ancestors[0].push(0);
    ancestors[0].push(1);
    ancestors[0].append(&mut vec);

    for k in 1..max_tree {
        for e in 0..n + 1 {
            let b_1 = ancestors[k - 1][e];
            let b_2 = ancestors[k - 1][b_1];
            ancestors[k].push(b_2);
        }
    }

    let mut depths: Vec<usize> = Vec::with_capacity(n + 1);
    depths.push(0);
    depths.push(1);
    for i in 2..ancestors[0].len() {
        depths.push(depths[ancestors[0][i]] + 1);
    }

    for _ in 0..q {
        let (mut a, mut b) = take_tuple();
        if depths[a] < depths[b] {
            std::mem::swap(&mut a, &mut b);
        }
        let d = depths[a] - depths[b];

        for i in 0..ancestors.len() {
            if (d >> i) & 1 == 1 {
                a = ancestors[i][a];
            }
        }

        if a == b {
            println!("{}", a);
            continue;
        }

        for i in (0..ancestors.len()).rev() {
            if ancestors[i][a] != ancestors[i][b] {
                a = ancestors[i][a];
                b = ancestors[i][b];
            }
        }
        println!("{}", ancestors[0][a]);
    }
}

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