graph/
lib.rs

1use std::mem;
2
3pub fn is_tree(n: usize, edges: &[(usize, usize)]) -> bool {
4    for &(a, b) in edges {
5        assert!(a < n);
6        assert!(b < n);
7        assert_ne!(a, b);
8    }
9
10    if n == 0 {
11        return true;
12    }
13
14    edges.len() == n - 1 && connectivity(n, edges)
15}
16
17pub fn connectivity(n: usize, edges: &[(usize, usize)]) -> bool {
18    fn dfs(i: usize, g: &[Vec<usize>], visited: &mut Vec<bool>) {
19        for &j in &g[i] {
20            if visited[j] {
21                continue;
22            }
23            visited[j] = true;
24            dfs(j, g, visited);
25        }
26    }
27
28    let mut g = vec![vec![]; n];
29    for &(a, b) in edges {
30        g[a].push(b);
31        g[b].push(a);
32    }
33    let mut visited = vec![false; n];
34    visited[0] = true;
35    dfs(0, &g, &mut visited);
36    visited.iter().filter(|&&f| f).count() == n
37}
38
39pub fn tree_drop_parent(
40    n: usize,
41    root: usize,
42    edges: &[(usize, usize)],
43) -> (Vec<Vec<usize>>, Vec<usize>) {
44    debug_assert!(is_tree(n, edges));
45
46    fn dfs(i: usize, p: usize, g: &Vec<Vec<usize>>, parent: &mut Vec<usize>) {
47        parent[i] = p;
48        for &j in &g[i] {
49            if j == p {
50                continue;
51            }
52            dfs(j, i, g, parent);
53        }
54    }
55
56    let mut g = vec![vec![]; n];
57    for &(a, b) in edges {
58        g[a].push(b);
59        g[b].push(a);
60    }
61    let mut parent = vec![usize::MAX; n];
62    dfs(root, root, &g, &mut parent);
63
64    for i in 0..n {
65        g[i] = mem::take(&mut g[i])
66            .into_iter()
67            .filter(|&j| j != parent[i])
68            .collect();
69    }
70
71    (g, parent)
72}
73
74#[cfg(test)]
75mod tests {
76    use crate::{is_tree, tree_drop_parent};
77
78    #[test]
79    fn test_is_tree_small() {
80        assert_eq!(is_tree(0, &[]), true);
81        assert_eq!(is_tree(1, &[]), true);
82        assert_eq!(is_tree(2, &[(0, 1)]), true);
83        assert_eq!(is_tree(3, &[(0, 1), (1, 2)]), true);
84        assert_eq!(is_tree(4, &[(0, 1), (0, 2), (0, 3)]), true);
85        assert_eq!(is_tree(4, &[(0, 1), (1, 2), (0, 3)]), true);
86        assert_eq!(is_tree(4, &[(0, 1), (2, 3)]), false);
87        assert_eq!(is_tree(4, &[(0, 1), (1, 2), (2, 0)]), false);
88    }
89
90    #[test]
91    fn test_drop_parent() {
92        assert_eq!(
93            // 0 (root) -- 1 -- 2 -- 3
94            tree_drop_parent(4, 0, &[(0, 1), (1, 2), (2, 3)]),
95            (
96                vec![vec![1], vec![2], vec![3], vec![]], // g
97                vec![0, 0, 1, 2],                        // parent
98            )
99        );
100    }
101}
graph/lib.rs

graph/
lib.rs
