The notebook of doubling.

abc_167.rs

use proconio::input;

fn main() {
    input! {
        n: usize,
        k: u64,
        a: [usize; n],
    }

    let result = solve(n, &a, k);
    println!("{}", result);
}

struct PopIter {
    index: usize,
    num: u64,
}

impl Iterator for PopIter {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        if self.num == 0 {
            return None;
        }
        let trailing_zeros = self.num.trailing_zeros();
        self.index += trailing_zeros as usize;
        self.num >>= trailing_zeros;
        let index = self.index;
        self.index += 1;
        self.num >>= 1;
        Some(index)
    }
}

fn solve(n: usize, a: &[usize], k: u64) -> u64 {
    let log2k = (k.next_power_of_two().trailing_zeros() + 1) as usize;

    // doubling[k][i] := the town stepped by k from i-th town
    let mut doubling = vec![vec![0; n]; log2k];
    for i in 0..n {
        // 1-indexed to 0-indexed
        doubling[0][i] = a[i] - 1;
    }
    for step in 1..log2k {
        for start in 0..n {
            doubling[step][start] = doubling[step - 1][doubling[step - 1][start]];
        }
    }

    let mut now = 0;
    for step in (PopIter { index: 0, num: k }) {
        now = doubling[step][now];
    }
    // 0-index to 1-indexed
    (now + 1) as u64
}

#[test]
fn case1() {
    assert_eq!(solve(4, &[3, 2, 4, 1], 5), 4);
}

#[test]
fn case2() {
    assert_eq!(solve(6, &[6, 5, 2, 5, 3, 5], 5727202214173249351), 2);
}