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

catlog/one/
tree_algorithms.rs
