AtCoder Beginner Contest 438 E

https://atcoder.jp/contests/abc438/tasks/abc438_e

バケツを渡した先を辿るとグラフになります。まずグラフを連結成分に分割します。各ノードの行先は一つなので、必ずループに陥ります。そうしたらあとは何回ループを回るかなどを考えればよいです。
まず、ループを作ります。インレットになっているノードがあればそこから辿って今まで出てきたことがあるノードまでを一つのブランチとします。水の量を出すときは、ノードがどのブランチに含まれるのかとそのブランチの何番目かを記憶しておいて、ブランチを辿って、最後の点がどのブランチに含まれるか調べて、というようにループに着くまで辿ります。こうすると意地悪な入力があるとO(N)回辿らないといけない場合も出てきますが、インレットの点の集合を作るときにHashSetを使うので、適当に混ぜられてそういう変なことにはならないと思われます。

// Heavy Buckets
#![allow(non_snake_case)]

use std::collections::{HashSet, HashMap};


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// UnionFind ////////////////////

use std::cmp::max;
type Node = usize;

struct UnionFind {
    parents: Vec<Node>,
    heights: Vec<i32>,
}

impl UnionFind {
    fn new(N: usize) -> UnionFind {
        let parents: Vec<Node> = (0..N).collect();
        let heights: Vec<i32> = vec![1; N];
        UnionFind { parents, heights }
    }
    
    fn len(&self) -> usize {
        self.parents.len()
    }
    
    fn is_root(&self, v: Node) -> bool {
        self.parents[v] == v
    }
    
    fn is_connected(&self, u: Node, v: Node) -> bool {
        self.root(u) == self.root(v)
    }
    
    fn connect(&mut self, u: Node, v: Node) {
        let r1 = self.root(u);
        let r2 = self.root(v);
        if r1 == r2 {
            return
        }
        
        let h1 = self.heights[r1];
        let h2 = self.heights[r2];
        if h1 <= h2 {   // r2にr1がぶら下がる
            self.parents[r1] = r2;
            self.heights[r2] = max(self.heights[r2], self.heights[r1]+1);
        }
        else {
            self.parents[r2] = r1;
            self.heights[r1] = max(self.heights[r1], self.heights[r2]+1);
        }
    }
    
    fn root(&self, v0: Node) -> Node {
        if self.is_root(v0) {
            v0
        }
        else {
            let p = self.parents[v0];
            self.root(p)
        }
    }
}


//////////////////// Branch ////////////////////

struct Branch {
    nodes: Vec<Node>,
    acc: Vec<usize>,
    dic: HashMap<Node, usize>   // ノードはブランチの何番目にあるか
}

impl Branch {
    fn len(&self) -> usize {
        self.nodes.len() - 1
    }
    
    fn is_loop(&self) -> bool {
        self.nodes[0] == self.end_point()
    }
    
    fn end_point(&self) -> Node {
        self.nodes[self.len()]
    }
    
    fn full_water(&self) -> usize {
        self.acc[self.len()]
    }
    
    // 指定した点Bから最大T回辿って、水の量と残り回数と端点を返す
    fn walk(&self, T: usize, B: Node) -> (usize, usize, Node) {
        let N = self.len();
        let i = *self.dic.get(&B).unwrap();
        if i+T < N {
            (self.acc[i+T] - self.acc[i], 0, 0)
        }
        else {
            (self.acc[N] - self.acc[i], i+T-N, self.end_point())
        }
    }
    
    fn find_point_on_loop(v0: Node, A: &Vec<usize>) -> Node {
        let mut visited: HashSet<Node> = HashSet::new();
        let mut v = v0;
        while !visited.contains(&v) {
            visited.insert(v);
            v = A[v]
        }
        v
    }
    
    fn create(v0: Node, A: &Vec<usize>, visited: &mut HashSet<Node>) -> Branch {
        let mut nodes: Vec<Node> = vec![v0];
        let mut acc: Vec<usize> = vec![0];
        let mut dic: HashMap<Node, usize> = HashMap::new();
        let mut v = v0;
        let mut i: usize = 0;
        while !visited.contains(&v) {
            dic.insert(v, i);
            visited.insert(v);
            acc.push(acc[i] + v + 1);
            v = A[v];
            i += 1;
            nodes.push(v);
        }
        if v != v0 {    // not loop
            acc.push(acc[i] + v + 1)
        }
        Branch { nodes, acc, dic }
    }
}


//////////////////// Graph ////////////////////

struct Graph {
    // 最初のPathがループ
    paths: Vec<Branch>,
    dic_paths: HashMap<Node, usize>     // ノードは何番目のpathに含まれるか
}

impl Graph {
    fn water(&self, T: usize, B: usize) -> usize {
        let mut s: usize = 0;
        let mut t = T;
        let mut v = B;
        loop {
            if let Some(&i) = self.dic_paths.get(&v) {
                let path = &self.paths[i];
                let (s1, t1, v1) = path.walk(t, v);
                s += s1;
                t = t1;
                v = v1;
                if t1 == 0 {
                    return s
                }
                if path.is_loop() {
                    break
                }
            }
        }
        
        // いま終点にいるはず
        let cycle = &self.paths[0];
        let q = t / cycle.len();
        let r = t % cycle.len();
        s += cycle.full_water() * q;
        s += cycle.acc[r] - cycle.acc[0];
        s
    }
    
    fn find_start_points(vs: &Vec<Node>, A: &Vec<usize>) -> Vec<Node> {
        let s1: HashSet<Node> = vs.iter().cloned().collect();
        let s2: HashSet<Node> = vs.iter().map(|&v| A[v]).collect();
        s1.difference(&s2).cloned().collect::<Vec<_>>()
    }
    
    fn create(vs: Vec<Node>, A: &Vec<usize>) -> Graph {
        let mut visited: HashSet<Node> = HashSet::new();
        let v0 = Branch::find_point_on_loop(vs[0], A);
        let cycle = Branch::create(v0, A, &mut visited);
        let mut paths: Vec<Branch> = vec![cycle];
        let mut dic_paths: HashMap<Node, usize> = HashMap::new();
        for &v in paths[0].nodes.iter() {
            dic_paths.insert(v, 0);
        }
        let start_points = Graph::find_start_points(&vs, A);
        for v0 in start_points {
            // ループ以外の部分
            let path = Branch::create(v0, A, &mut visited);
            for i in 0..path.len() {
                dic_paths.insert(path.nodes[i], paths.len());
            }
            paths.push(path);
        }
        Graph { paths, dic_paths }
    }
}


//////////////////// process ////////////////////

fn read_input() -> (usize, Vec<usize>) {
    let v: Vec<usize> = read_vec();
    let Q = v[1];
    let A: Vec<usize> = read_vec::<usize>().into_iter().map(|a| a-1).collect();
    (Q, A)
}

fn divide_into_connected(uf: &UnionFind) -> HashMap<usize, Vec<usize>> {
    let N = uf.len();
    let mut m: HashMap<usize, Vec<usize>> = HashMap::new();
    for i in 0..N {
        let r = uf.root(i);
        let e = m.entry(r).or_insert(vec![]);
        (*e).push(i)
    }
    m
}

type Query = (usize, Node);

fn read_query() -> Query {
    let v: Vec<usize> = read_vec();
    let (T, B) = (v[0], v[1]-1);
    (T, B)
}

fn F(Q: usize, A: Vec<usize>) {
    let qs: Vec<Query> = (0..Q).map(|_| read_query()).collect();
    let N = A.len();
    let mut uf = UnionFind::new(N);
    for i in 0..N {
        uf.connect(i, A[i])
    }
    
    let subgroups = divide_into_connected(&uf);
    let graphs: HashMap<Node, Graph> =
                                subgroups.into_iter().
                                map(|(r, subg)| (r, Graph::create(subg, &A))).
                                collect();
    
    for (T, B) in qs {
        let root = uf.root(B);
        if let Some(graph) = graphs.get(&root) {
            println!("{}", graph.water(T, B))
        }
    }
}

fn main() {
    let (Q, A) = read_input();
    F(Q, A)
}