catlog/one/
tree_algorithms.rs

1//! Algorithms on trees.
2
3use ego_tree::iter::{Descendants, Edge};
4use ego_tree::{NodeId, NodeRef, Tree};
5use itertools::{EitherOrBoth::Both, Itertools};
6use std::collections::VecDeque;
7
8/// Extension trait adding traversal algorithms on [trees](Tree).
9pub trait TreeTraversal<T> {
10    /// Iterates over descendants of node in depth-first order.
11    fn dfs(&self) -> Descendants<'_, T>;
12
13    /// Iterates over descendants of node in breadth-first order.
14    fn bfs(&self) -> BreadthFirstTraversal<'_, T>;
15
16    /// Iterates over left boundary of node.
17    fn left_boundary(&self) -> impl Iterator<Item = Self>;
18
19    /// Iterates over right boundary of node.
20    fn right_boundary(&self) -> impl Iterator<Item = Self>;
21}
22
23/// Iterator for traversing a tree in breadth-first order.
24pub struct BreadthFirstTraversal<'a, T: 'a> {
25    tree: &'a Tree<T>,
26    queue: VecDeque<(NodeId, usize)>,
27    current_level: usize,
28}
29
30impl<'a, T: 'a> BreadthFirstTraversal<'a, T> {
31    /// Initialize a breadth-first traversal at the given node.
32    pub fn starting_at(root: NodeRef<'a, T>) -> Self {
33        let tree = root.tree();
34        let mut queue = VecDeque::new();
35        queue.push_back((root.id(), 1));
36        Self { tree, queue, current_level: 0 }
37    }
38
39    /// Peeks at the next node, if it's at the same level as the previous one.
40    pub fn peek_at_same_level(&self) -> Option<NodeRef<'a, T>> {
41        self.queue.front().and_then(|(id, level)| {
42            if *level == self.current_level {
43                self.tree.get(*id)
44            } else {
45                None
46            }
47        })
48    }
49}
50
51impl<'a, T: 'a> Iterator for BreadthFirstTraversal<'a, T> {
52    type Item = NodeRef<'a, T>;
53
54    fn next(&mut self) -> Option<Self::Item> {
55        let (id, level) = self.queue.pop_front()?;
56        self.current_level = level;
57        let node = self.tree.get(id).unwrap();
58        for child in node.children() {
59            self.queue.push_back((child.id(), level + 1));
60        }
61        Some(node)
62    }
63}
64
65impl<'a, T: 'a> std::iter::FusedIterator for BreadthFirstTraversal<'a, T> {}
66
67impl<'a, T: 'a> TreeTraversal<T> for NodeRef<'a, T> {
68    /// Uses the built-in traversal algorithm, which is depth-first, though that
69    /// is not documented: <https://github.com/rust-scraper/ego-tree/issues/38>.
70    fn dfs(&self) -> Descendants<'a, T> {
71        self.descendants()
72    }
73
74    /// Implements the standard BFS algorithm using a queue.
75    fn bfs(&self) -> BreadthFirstTraversal<'a, T> {
76        BreadthFirstTraversal::starting_at(*self)
77    }
78
79    fn left_boundary(&self) -> impl Iterator<Item = Self> {
80        let mut maybe_node = Some(*self);
81        std::iter::from_fn(move || {
82            let prev = maybe_node;
83            maybe_node = maybe_node.and_then(|node| node.first_child());
84            prev
85        })
86    }
87
88    fn right_boundary(&self) -> impl Iterator<Item = Self> {
89        let mut maybe_node = Some(*self);
90        std::iter::from_fn(move || {
91            let prev = maybe_node;
92            maybe_node = maybe_node.and_then(|node| node.last_child());
93            prev
94        })
95    }
96}
97
98/// Extension trait adding isomorphism checking on [trees](Tree).
99pub trait TreeIsomorphism<T> {
100    /// Is the tree isomorphic to another?
101    ///
102    /// The standard data structure for trees based on pointers has only one notion
103    /// of "sameness" that makes sense, but for vector-backed trees with node IDs,
104    /// trees can be isomorphic (logically the same) without having underlying data
105    /// that is equal. This methods checks for logical sameness.
106    ///
107    /// Note that the isomorphism check ignores orphaned nodes, since those are
108    /// logically deleted.
109    fn is_isomorphic_to(&self, other: &Self) -> bool;
110}
111
112impl<T> TreeIsomorphism<T> for Tree<T>
113where
114    T: Eq,
115{
116    fn is_isomorphic_to(&self, other: &Self) -> bool {
117        self.root()
118            .traverse()
119            .zip_longest(other.root().traverse())
120            .all(|pair| match pair {
121                Both(Edge::Open(n1), Edge::Open(n2)) | Both(Edge::Close(n1), Edge::Close(n2)) => {
122                    n1.value() == n2.value()
123                }
124                _ => false,
125            })
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use ego_tree::tree;
133
134    #[test]
135    fn dfs() {
136        let tree = tree!('a' => { 'b' => { 'd', 'e' }, 'c' });
137        let values: Vec<_> = tree.root().dfs().map(|node| *node.value()).collect();
138        assert_eq!(values, vec!['a', 'b', 'd', 'e', 'c']);
139    }
140
141    #[test]
142    fn bfs() {
143        let tree = tree!('a' => { 'b' => { 'd', 'e' }, 'c' });
144        let values: Vec<_> = tree.root().bfs().map(|node| *node.value()).collect();
145        assert_eq!(values, vec!['a', 'b', 'c', 'd', 'e']);
146
147        let tree = tree!('a' => { 'b' => {'d'}, 'c' => {'e'} });
148        let root = tree.root();
149        let mut traverse = root.bfs();
150        traverse.next();
151        assert!(traverse.peek_at_same_level().is_none());
152        assert_eq!(traverse.nth(2).map(|node| *node.value()), Some('d'));
153        assert_eq!(traverse.peek_at_same_level().map(|node| *node.value()), Some('e'));
154    }
155
156    #[test]
157    fn isomorphism() {
158        let tree = tree!('a' => { 'b' => { 'd', 'e' }, 'c' });
159        assert!(tree.is_isomorphic_to(&tree));
160
161        let other = tree!('a' => { 'b' => { 'd' }, 'e' => { 'c' }});
162        let tree_dfs_values: Vec<_> = tree.root().dfs().map(|node| *node.value()).collect();
163        let other_dfs_values: Vec<_> = other.root().dfs().map(|node| *node.value()).collect();
164        assert_eq!(tree_dfs_values, other_dfs_values);
165        assert!(!tree.is_isomorphic_to(&other));
166    }
167}
catlog/one/tree_algorithms.rs

catlog/one/
tree_algorithms.rs
