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
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! Often when working with trees we are given them as a graph with undirected edges, however
//! sometimes a better representation is a rooted tree.
//!
//! - Time Complexity: O(V+E)
//!
//! # Resources
//!
//! - [W. Fiset's video](https://www.youtube.com/watch?v=2FFq2_je7Lg&list=PLDV1Zeh2NRsDGO4--qE8yH72HFL1Km93P&index=9)

use crate::algo::graph::UnweightedAdjacencyList;

#[derive(Debug, Eq, PartialEq)]
pub struct TreeNode {
    pub id: usize,
    pub children: Vec<TreeNode>,
}

impl TreeNode {
    pub fn new(id: usize) -> Self {
        Self {
            id,
            children: vec![],
        }
    }
    pub fn from_adjacency_list(graph: &UnweightedAdjacencyList, root: usize) -> Self {
        fn build_tree_recursive(
            graph: &UnweightedAdjacencyList,
            node_id: usize,
            parent_id: Option<usize>,
        ) -> TreeNode {
            let mut node = TreeNode::new(node_id);
            for &child_id in &graph[node_id] {
                if let Some(id) = parent_id {
                    if id == child_id {
                        continue;
                    }
                }
                let child_node = build_tree_recursive(graph, child_id, Some(node_id));
                node.children.push(child_node);
            }
            node
        }
        build_tree_recursive(graph, root, None)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_tree_rooting() {
        let mut graph = UnweightedAdjacencyList::with_size(9);
        graph.add_undirected_edge(0, 1);
        graph.add_undirected_edge(2, 1);
        graph.add_undirected_edge(2, 3);
        graph.add_undirected_edge(3, 4);
        graph.add_undirected_edge(5, 3);
        graph.add_undirected_edge(2, 6);
        graph.add_undirected_edge(6, 7);
        graph.add_undirected_edge(6, 8);
        let tree = TreeNode::from_adjacency_list(&graph, 6);
        // Rooted at 6 the tree should look like:
        //         6
        //      2  7  8
        //    1  3
        //  0   4 5
        println!("{:?}", &tree);
        assert_eq!(
            tree,
            TreeNode {
                id: 6,
                children: vec![
                    TreeNode {
                        id: 2,
                        children: vec![
                            TreeNode {
                                id: 1,
                                children: vec![TreeNode::new(0)]
                            },
                            TreeNode {
                                id: 3,
                                children: vec![TreeNode::new(4), TreeNode::new(5)]
                            }
                        ]
                    },
                    TreeNode::new(7),
                    TreeNode::new(8)
                ]
            }
        );
        let tree = TreeNode::from_adjacency_list(&graph, 3);
        // Rooted at 3 the tree should look like:
        //       3
        //    2  4  5
        //  6  1
        // 7 8  0
        println!("{:?}", &tree);
        assert_eq!(
            tree,
            TreeNode {
                id: 3,
                children: vec![
                    TreeNode {
                        id: 2,
                        children: vec![
                            TreeNode {
                                id: 1,
                                children: vec![TreeNode::new(0)]
                            },
                            TreeNode {
                                id: 6,
                                children: vec![TreeNode::new(7), TreeNode::new(8)]
                            }
                        ]
                    },
                    TreeNode::new(4),
                    TreeNode::new(5)
                ]
            }
        );
    }
}