Skip to content

Instantly share code, notes, and snippets.

@takuma-saito
Last active May 21, 2023 10:04
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save takuma-saito/0cd0cffade3292e25d1909314cfe061a to your computer and use it in GitHub Desktop.
Save takuma-saito/0cd0cffade3292e25d1909314cfe061a to your computer and use it in GitHub Desktop.
struct TreeNode<T> {
value: T,
left: TreeNodeRef<T>,
right: TreeNodeRef<T>
}
type TreeNodeRef<T> = Option<Box<TreeNode<T>>>;
type Node = TreeNode<i64>;
type NodeRef = TreeNodeRef<i64>;
fn make_node(v: i64) -> Node {
Node { value: v, left: None, right: None }
}
fn make_node_ref(node: Node) -> NodeRef {
Some(Box::new(node))
}
fn make_node_leaf(v: i64) -> NodeRef {
make_node_ref(make_node(v))
}
// left -> right
struct TreeNodeIterator<'a, T> {
unvisited: Vec<(&'a TreeNode<T>, usize)>
}
impl<'a, T> Iterator for TreeNodeIterator<'a, T> {
type Item = (&'a TreeNode<T>, usize);
fn next(&mut self) -> Option<(&'a TreeNode<T>, usize)> {
if let Some((node, depth)) = self.unvisited.pop() {
let left = &node.left;
let right = &node.right;
if let Some(n) = right { self.unvisited.push((&*n, depth+1)); }
if let Some(n) = left { self.unvisited.push((&*n, depth+1)); }
Some((node, depth)) // preorder
} else { None }
}
}
impl<'a, T> IntoIterator for &'a TreeNode<T> {
type Item = (&'a TreeNode<T>, usize);
type IntoIter = TreeNodeIterator<'a, T>;
fn into_iter(self) -> TreeNodeIterator<'a, T> {
TreeNodeIterator {
unvisited: vec![(self, 0)]
}
}
}
fn main() {
let tree = Node {
value: 10,
left: make_node_leaf(5),
right: make_node_ref(Node {
value: 3,
left: None,
right: make_node_leaf(9)
})
};
for (node, depth) in &tree {
println!("{}{}", " ".repeat(4*depth), node.value);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment