Submission #56875669

---

Source Code Expand

use std::iter::successors;

use nekolib::fmt::SpaceSep;
use proconio::{input, marker::Usize1};

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

    let f = N1FnGraph::new(&x);
    let xk: Vec<_> = (0..n).map(|i| f.step(i, k)).collect();
    let res: Vec<_> = xk.iter().map(|&i| a[i]).collect();
    println!("{}", SpaceSep(&res));
}

struct N1FnGraph {
    la: N1La,
    cycle: Vec<Vec<usize>>,
}

impl N1FnGraph {
    pub fn new(f: &[usize]) -> Self {
        let (par, cycle) = Self::par_cycle(f);
        let la = N1La::new(par);
        Self { la, cycle }
    }
    pub fn step(&self, v: usize, k: u64) -> usize {
        let d = self.la.depth(v);
        debug_assert!(d > 0); // for sentinel `n`
        if ((d - 1) as u64) >= k {
            self.la.ascend(v, k as _)
        } else {
            let r = self.la.ascend(v, d - 1);
            let k = k - (d - 1) as u64;
            let lambda = self.cycle[r].len();
            let i = (k % lambda as u64) as usize;
            self.cycle[r][i]
        }
    }

    fn par_cycle(f: &[usize]) -> (Vec<usize>, Vec<Vec<usize>>) {
        let n = f.len();
        let mut par = f.to_owned();
        par.push(n + 1);
        let mut cycle = vec![vec![]; n];

        let mut seen = vec![false; n];
        let mut stack = vec![];
        let mut index = vec![n; n];
        for mut i in 0..n {
            while !seen[i] {
                if index[i] < n {
                    break;
                }
                index[i] = stack.len();
                stack.push(i);
                i = f[i];
            }
            if !seen[i] {
                let s = i;
                let mut c = vec![];
                while let Some(j) = stack.pop() {
                    c.push(j);
                    seen[j] = true;
                    if j == s {
                        break;
                    }
                }
                let r = c[0];
                c.reverse();
                c.rotate_right(1);
                cycle[r] = c;
                par[r] = n;
            }
            while let Some(j) = stack.pop() {
                seen[j] = true;
            }
        }
        (par, cycle)
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum RefNode {
    MicroRoot(usize, usize),
    MacroLeaf(usize),
}
use RefNode::{MacroLeaf, MicroRoot};

pub struct N1La {
    par: Vec<usize>,
    depth: Vec<usize>,
    ref_node: Vec<RefNode>,
    micro_tree: Vec<Option<(usize, Vec<usize>)>>,
    jump_pointer: Vec<Vec<usize>>,
    ladder: Vec<Vec<usize>>,
    which_ladder: Vec<(usize, usize)>,
    table: Vec<Vec<Vec<usize>>>,
}

impl N1La {
    pub fn new(par: Vec<usize>) -> Self {
        let n = par.len();
        let thold = Self::thold(n);
        let (g, r) = Self::children(&par);
        let (depth, size) = Self::depth_size(&g, r);
        let ref_node = Self::ref_node(&g, r, &depth, &size, &par, thold);
        let micro_tree = Self::micro_tree(&g, &ref_node);
        let jump_pointer = Self::jump_pointer(&g, r, &ref_node);
        let (ladder, which_ladder) = Self::ladder(&g, &ref_node, &par);
        let table = Self::table(&micro_tree, thold);
        Self {
            par,
            depth,
            ref_node,
            micro_tree,
            jump_pointer,
            ladder,
            which_ladder,
            table,
        }
    }
    pub fn la(&self, v: usize, l: usize) -> usize {
        assert!(self.depth[v] >= l);
        self.ascend(v, self.depth(v) - l)
    }
    pub fn ascend(&self, v: usize, a: usize) -> usize {
        assert!(self.depth[v] >= a);
        let (v, a) = match self.ref_node[v] {
            MicroRoot(r, ord) => {
                if self.depth[v] - self.depth[r] >= a {
                    let &(e, ref vs) = self.micro_tree[r].as_ref().unwrap();
                    return vs[self.lookup(e, ord, a)];
                } else {
                    (self.par[r], a - (self.depth[v] - self.depth[r] + 1))
                }
            }
            _ => (v, a),
        };
        let (v, a) = match self.ref_node[v] {
            MacroLeaf(u) => (u, a + (self.depth[u] - self.depth[v])),
            MicroRoot(..) => unreachable!(),
        };
        // now v is a macro leaf, so it has jump pointers.

        if a == 0 {
            return v;
        }

        // use the appropriate jump pointer
        let (v, a) = {
            let msb = (usize::BITS - 1 - a.leading_zeros()) as usize;
            (self.jump_pointer[v][msb], a - (1 << msb))
        };

        // use the appropriate ladder
        let (i, j) = self.which_ladder[v];
        self.ladder[i][a + j]
    }
    pub fn depth(&self, v: usize) -> usize { self.depth[v] }

    fn thold(n: usize) -> usize {
        (1..).take_while(|&i| (1_usize << (4 * i)) <= n).last().unwrap_or(1)
    }
    fn children(par: &[usize]) -> (Vec<Vec<usize>>, usize) {
        let n = par.len();
        let r = (0..n).find(|&i| par[i] == n).unwrap();
        let mut g = vec![vec![]; n];
        for i in 0..n {
            if par[i] < n {
                g[par[i]].push(i);
            }
        }
        (g, r)
    }
    fn depth_size(g: &[Vec<usize>], r: usize) -> (Vec<usize>, Vec<usize>) {
        fn dfs(
            g: &[Vec<usize>],
            v: usize,
            depth: &mut [usize],
            size: &mut [usize],
        ) {
            for &nv in &g[v] {
                depth[nv] = depth[v] + 1;
                dfs(g, nv, depth, size);
                size[v] += size[nv];
            }
        }
        let n = g.len();
        let mut depth = vec![0; n];
        let mut size = vec![1; n];
        dfs(g, r, &mut depth, &mut size);
        (depth, size)
    }
    fn ref_node(
        g: &[Vec<usize>],
        r: usize,
        depth: &[usize],
        size: &[usize],
        par: &[usize],
        thold: usize,
    ) -> Vec<RefNode> {
        fn dfs(
            g: &[Vec<usize>],
            v: usize,
            size: &[usize],
            thold: usize,
            res: &mut [RefNode],
            micro: &mut Option<&mut Vec<usize>>,
        ) {
            if let Some(micro) = micro.as_mut() {
                micro.push(v);
            }
            for &nv in &g[v] {
                if size[nv] <= thold && size[v] > thold {
                    debug_assert!(micro.is_none());
                    let mut micro = vec![];
                    dfs(g, nv, size, thold, res, &mut Some(&mut micro));
                    for (i, &u) in micro.iter().enumerate() {
                        res[u] = MicroRoot(nv, i);
                    }
                } else {
                    dfs(g, nv, size, thold, res, micro);
                }
            }
        }
        let n = g.len();
        let mut res = vec![MacroLeaf(n); n];
        dfs(g, r, size, thold, &mut res, &mut None);

        let mut by_d = vec![vec![]; n];
        for v in 0..n {
            by_d[depth[v]].push(v);
        }
        for d in (0..n).rev() {
            for &v in &by_d[d] {
                if matches!(res[v], MicroRoot(..)) {
                    continue;
                }
                for u in successors(Some(v), |&v| (v < n).then(|| par[v]))
                    .take_while(|&u| u < n)
                {
                    if res[u] != MacroLeaf(n) {
                        break;
                    }
                    res[u] = MacroLeaf(v);
                }
            }
        }

        res
    }
    fn micro_tree(
        g: &[Vec<usize>],
        ref_node: &[RefNode],
    ) -> Vec<Option<(usize, Vec<usize>)>> {
        fn dfs(
            g: &[Vec<usize>],
            v: usize,
            enc: &mut usize,
            vs: &mut Vec<usize>,
        ) {
            *enc = *enc << 1 | 1;
            vs.push(v);
            for &nv in &g[v] {
                dfs(g, nv, enc, vs);
            }
            *enc = *enc << 1 | 0;
        }
        let n = g.len();
        let mut res = vec![None; n];
        for v in 0..n {
            if let MicroRoot(_, 0) = ref_node[v] {
                let mut enc = 0;
                let mut vs = vec![];
                dfs(g, v, &mut enc, &mut vs);
                res[v] = Some((enc, vs));
            }
        }
        res
    }
    fn jump_pointer(
        g: &[Vec<usize>],
        r: usize,
        ref_node: &[RefNode],
    ) -> Vec<Vec<usize>> {
        fn dfs(
            g: &[Vec<usize>],
            v: usize,
            ref_node: &[RefNode],
            path: &mut Vec<usize>,
            res: &mut [Vec<usize>],
        ) {
            path.push(v);
            if let MacroLeaf(l) = ref_node[v] {
                if v == l {
                    let pl = path.len();
                    res[v] = (0..)
                        .map(|i| 1 << i)
                        .take_while(|&i| i < pl)
                        .map(|i| path[pl - (i + 1)])
                        .collect();
                } else {
                    for &nv in &g[v] {
                        dfs(g, nv, ref_node, path, res);
                    }
                }
            }
            path.pop();
        }

        let n = g.len();
        let mut res = vec![vec![]; n];
        dfs(g, r, ref_node, &mut vec![], &mut res);
        res
    }
    fn ladder(
        g: &[Vec<usize>],
        ref_node: &[RefNode],
        par: &[usize],
    ) -> (Vec<Vec<usize>>, Vec<(usize, usize)>) {
        let n = g.len();

        let mut ladder = vec![vec![]; n];
        let mut which_ladder = vec![(n, n); n];
        let next_par1 = |v: usize| {
            (par[v] < n && ref_node[par[v]] == ref_node[v]).then(|| par[v])
        };
        let next_par2 = |v: usize| (par[v] < n).then(|| par[v]);
        for v in 0..n {
            match ref_node[v] {
                MacroLeaf(r) if r == v => (),
                _ => continue,
            }
            ladder[v] = successors(Some(v), |&v| next_par1(v)).collect();
            for (i, &u) in ladder[v].iter().enumerate() {
                which_ladder[u] = (v, i);
            }

            let u = par[*ladder[v].last().unwrap()];
            let k = ladder[v].len();
            ladder[v].extend(
                successors((u < n).then(|| u), |&v| next_par2(v)).take(k),
            );
        }
        (ladder, which_ladder)
    }
    fn table(
        micro_tree: &[Option<(usize, Vec<usize>)>],
        thold: usize,
    ) -> Vec<Vec<Vec<usize>>> {
        let mut res = vec![vec![]; 1 << (2 * thold)];
        for &(e, ref vs) in micro_tree.iter().filter_map(|o| o.as_ref()) {
            let len = vs.len();
            let mut v = 0;
            let mut path = vec![];
            let mut res_e = vec![vec![]; len];
            for b in (0..2 * len).rev().map(|i| e >> i & 1 != 0) {
                if b {
                    path.push(v);
                    for &u in path.iter().rev() {
                        res_e[v].push(u);
                    }
                    v += 1;
                } else {
                    path.pop();
                }
            }
            res[e] = res_e;
        }
        res
    }

    fn lookup(&self, e: usize, v: usize, a: usize) -> usize {
        self.table[e][v][a]
    }
}

/// This module is bundled automatically.
/// See <https://rsk0315.github.io/nekolib/nekolib_doc/index.html> for documentation.
/// Commit: d829603bb021542857778d68b3de948ff48aff91-dirty
#[allow(unused)]
pub mod nekolib {
    pub mod fmt {
        pub mod str_sep {
            use std::fmt;
            pub struct SpaceSep<I>(pub I);
            pub struct PerLine<I>(pub I);
            pub struct StrSep<'a, I>(pub I, pub &'a str);
            pub struct SpaceSepUsize1<I>(pub I);
            pub struct PerLineUsize1<I>(pub I);
            pub struct StrSepUsize1<'a, I>(pub I, pub &'a str);
            macro_rules! impl_fmt {
                ( $( $fmt:ident )* ) => { $(
                    #[allow(non_snake_case)]
                    fn $fmt<I, T: fmt::$fmt>(iter: I, sep: &str, f: &mut fmt::Formatter<'_>) -> fmt::Result
                    where
                        I: IntoIterator<Item = T>,
                    {
                        let mut iter = iter.into_iter();
                        if let Some(first) = iter.by_ref().next() {
                            first.fmt(f)?;
                        }
                        iter.map(|rest| { write!(f, "{}", sep)?; rest.fmt(f) }).collect()
                    }

                    impl<I, T: fmt::$fmt> fmt::$fmt for SpaceSep<I>
                    where
                        I: IntoIterator<Item = T> + Clone,
                    {
                        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                            $fmt(self.0.clone(), " ", f)
                        }
                    }
                    impl<I, T: fmt::$fmt> fmt::$fmt for PerLine<I>
                    where
                        I: IntoIterator<Item = T> + Clone,
                    {
                        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                            $fmt(self.0.clone(), "\n", f)
                        }
                    }
                    impl<I, T: fmt::$fmt> fmt::$fmt for StrSep<'_, I>
                    where
                        I: IntoIterator<Item = T> + Clone,
                    {
                        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                            $fmt(self.0.clone(), self.1, f)
                        }
                    }
                )* }
            }
            macro_rules! impl_fmt_usize1 {
                ( $( $fmt:ident )* ) => { $(
                    impl<I, T> fmt::$fmt for SpaceSepUsize1<I>
                    where
                        I: IntoIterator<Item = T> + Clone,
                        T: std::ops::Add<usize, Output = usize>,
                    {
                        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                            $fmt(self.0.clone().into_iter().map(|u| u + 1), " ", f)
                        }
                    }
                    impl<I, T> fmt::$fmt for PerLineUsize1<I>
                    where
                        I: IntoIterator<Item = T> + Clone,
                        T: std::ops::Add<usize, Output = usize>,
                    {
                        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                            $fmt(self.0.clone().into_iter().map(|u| u + 1), "\n", f)
                        }
                    }
                    impl<I, T> fmt::$fmt for StrSepUsize1<'_, I>
                    where
                        I: IntoIterator<Item = T> + Clone,
                        T: std::ops::Add<usize, Output = usize>,
                    {
                        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                            $fmt(self.0.clone().into_iter().map(|u| u + 1), self.1, f)
                        }
                    }
            )* }
            }
            impl_fmt! { Binary Debug Display LowerExp LowerHex Octal Pointer UpperExp UpperHex }
            impl_fmt_usize1! { Debug Display LowerHex Octal UpperHex }
        }
        #[allow(unused_imports)]
        pub use str_sep::*;
    }
}

Submission Info

Judge Result