Skip to main content

library_checker/tree/
jump_on_tree.rs

1use competitive::graph::TreeGraphScanner;
2use competitive::prelude::*;
3
4#[verify::library_checker("jump_on_tree")]
5pub fn jump_on_tree(reader: impl Read, writer: impl Write) {
6    prepare_io!(reader, writer);
7    sc!(n, q, (g, _): @TreeGraphScanner::<usize>::new(n));
8    let hld = g.hld(0);
9    for _ in 0..q {
10        sc!(s, t, i);
11        pp!(hld.jump(s, t, i).unwrap_or(!0) as isize);
12    }
13}
14
15#[verify::library_checker("jump_on_tree")]
16pub fn jump_on_tree_level_ancestor(reader: impl Read, writer: impl Write) {
17    prepare_io!(reader, writer);
18    sc!(n, q, (g, _): @TreeGraphScanner::<usize>::new(n));
19    let la = g.level_ancestor(0);
20    let lca = g.lca(0);
21    for _ in 0..q {
22        sc!(s, t, i);
23        let l = lca.lca(s, t);
24        let dl = la.depth(l);
25        let ds = la.depth(s) - dl;
26        let dt = la.depth(t) - dl;
27        let ans = if i <= ds {
28            la.la(s, i)
29        } else if i <= ds + dt {
30            la.la(t, ds + dt - i)
31        } else {
32            None
33        };
34        pp!(ans.unwrap_or(!0) as isize);
35    }
36}
37
38#[verify::library_checker("jump_on_tree")]
39pub fn jump_on_tree_level_ancestor_batch(reader: impl Read, writer: impl Write) {
40    prepare_io!(reader, writer);
41    sc!(n, q, (g, _): @TreeGraphScanner::<usize>::new(n), queries: [(usize, usize, usize); iter q]);
42    let lca = g.lca(0);
43    let results = g.level_ancestor_batch(
44        0,
45        queries.map(|(s, t, i)| {
46            let l = lca.lca(s, t);
47            let dl = lca.depth(l);
48            let ds = lca.depth(s) - dl;
49            let dt = lca.depth(t) - dl;
50            if i <= ds {
51                (s, i)
52            } else if i <= ds + dt {
53                (t, ds + dt - i)
54            } else {
55                (0, n)
56            }
57        }),
58    );
59    pp!(@lf @it results.iter().map(|&v| v.unwrap_or(!0) as isize));
60}