Submission #56852439

---

Source Code Expand

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 = f.step_each(k);
    let res: Vec<_> = xk.iter().map(|&i| a[i]).collect();
    println!("{}", SpaceSep(&res));
}

struct N1FnGraph {
    g: Vec<Vec<usize>>,
    cycle: Vec<Vec<usize>>,
}

impl N1FnGraph {
    pub fn new(f: &[usize]) -> Self {
        let n = f.len();
        let (par, cycle) = Self::par_cycle(f);
        let mut g = vec![vec![]; n + 1];
        for i in 0..n {
            g[par[i]].push(i);
        }
        Self { g, cycle }
    }

    pub fn step_each(&self, k: u64) -> Vec<usize> {
        fn dfs(
            g: &[Vec<usize>],
            v: usize,
            path: &mut Vec<usize>,
            k: u64,
            res: &mut [usize],
            cycle: &[Vec<usize>],
        ) {
            if v < res.len() {
                // path: [n, ..., v]
                let pl = path.len();
                res[v] = if (pl as u64) - 1 > k {
                    path[pl - 1 - k as usize]
                } else {
                    let r = path[1];
                    let cl = cycle[r].len();
                    let i = (k - (pl - 1) as u64) % cl as u64;
                    cycle[r][i as usize]
                };
            }

            for &nv in &g[v] {
                path.push(nv);
                dfs(g, nv, path, k, res, cycle);
                path.pop();
            }
        }
        let n = self.g.len() - 1;
        let mut res = vec![n; n];
        dfs(&self.g, n, &mut vec![n], k, &mut res, &self.cycle);
        res
    }

    fn par_cycle(f: &[usize]) -> (Vec<usize>, Vec<Vec<usize>>) {
        let n = f.len();
        let mut par = f.to_owned();
        par.push(n);
        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();
                cycle[r] = c;
                par[r] = n;
            }
            while let Some(j) = stack.pop() {
                seen[j] = true;
            }
        }
        (par, cycle)
    }
}

/// 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