dsalgo 0.3.10

A package for Datastructures and Algorithms.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

//! find a node step by k edges from u for each query in offline

use crate::{
    level_ancestor_offline_with_dfs_recurse::level_ancestor,
    tree_diameter_terminal_nodes::diameter_ends,
};

pub fn kth_nodes(
    g: &[Vec<usize>],
    queries: &[(usize, usize)],
) -> Vec<Option<usize>> {
    let (u, v) = diameter_ends(g);

    let res0 = level_ancestor(g, u, queries);

    let res1 = level_ancestor(g, v, queries);

    res0.into_iter()
        .zip(res1.into_iter())
        .map(|(u, v)| {
            if u.is_some() {
                u
            } else if v.is_some() {
                v
            } else {
                None
            }
        })
        .collect()
}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test() {
        let cases = vec![(
            5,
            vec![(1, 2), (2, 3), (3, 4), (3, 5)],
            vec![(2, 2), (5, 3), (3, 3)],
            vec![4, 1, -1],
        )];

        use crate::{
            adjacency_list_graph_from_edges::graph_from_edges,
            edges_to_0_indexed::edges_to_0_indexed,
        };

        for (n, edges, mut queries, ans) in cases {
            let g = graph_from_edges(n, &edges_to_0_indexed(edges), false);

            for q in queries.iter_mut() {
                q.0 = q.0 - 1;
            }

            let res = kth_nodes(&g, &queries);

            for i in 0..res.len() {
                let x =
                    if ans[i] == -1 { None } else { Some(ans[i] as usize - 1) };

                assert_eq!(res[i], x);
            }
        }
    }
}