Problem overview
A binary tree was written as a string by walking it in preorder: a node first, then its whole left subtree, then its whole right subtree. Before each node's value the walk writes one dash per level of depth, so the root has none, its children have one, their children two, and so on. When a node has a single child, that child is the left one.
The function recover_from_preorder receives that string, traversal, and must rebuild the tree and return its root as an Option<Box<TreeNode>>.
For "1-2--3--4-5--6--7" the nodes are 1 at depth 0, then 2 at depth 1, 3 and 4 at depth 2, then 5 at depth 1, and 6 and 7 at depth 2. The tree has 1 as the root, 2 and 5 as its children, 3 and 4 under 2, and 6 and 7 under 5. Values can have several digits, as in "5-12--7---40--9": the dashes before a value end where its digits start.
Approach: keep the current path on a stack
Because the string is a preorder walk, the node being read always belongs somewhere along the path from the root to the node read just before it. The node's depth tells exactly where: its parent is the node at depth depth - 1 on that path.
So the code keeps a stack of the nodes on the current path, with the root at the bottom. The node at position d of the stack has depth d. For each new node it does two things:
- Pops the stack down to
depthnodes. The popped nodes are finished: everything after them in the string is shallower, so their subtrees are complete. - Pushes the new node. The node now on top of the stack is its parent.
In Rust there is one wrinkle. A Box has a single owner, so a node can't sit in the stack and in its parent's left field at the same time. The solution makes the stack own the nodes of the path and moves a node into its parent only when it is popped. A parent's first finished child becomes its left child and the second becomes its right child, which also puts a single child on the left, as the statement requires.
Take "1-2--3---4-5--6---7". After reading 4 the stack is 1, 2, 3, 4. The next node, 5, has depth 1, so the stack must shrink to one node. Popping 4 attaches it to 3, popping 3 attaches it to 2, and popping 2 attaches it to 1. Then 5 is pushed. The tree ends up as [1, 2, 5, 3, null, 6, null, 4, null, 7], with 3 and 6 as left children because they are the only child of their parents.
Step by step
- Read the string by index.
traversal.as_bytes()gives a byte slice. The string holds only digits and dashes, one byte each, sobytes[i]is the character at positioni. - Create the stack as an empty
Vec<Box<TreeNode>>. - Create the position
iat 0. Every loop that reads a character advances it, so the string is read once. - Outer loop: run while
i < bytes.len(). Each pass reads exactly one node, its dashes and then its digits. - Count the dashes.
depthstarts at 0, and an inner loop adds 1 and movesiforward whilebytes[i]isb'-'. It needs no bounds check, because a value always follows the dashes. - Read the value.
valuestarts at 0. Whileiis inside the string and the byte isn't a dash, multiplyvalueby 10 and add the digit. Reading 4 then 0 gives 4, then 40. - Fix the stack. Call the helper
Self::pop_to(&mut stack, depth)so that only the new node's ancestors remain. - Push the new node, a
BoxaroundTreeNode::new(value). - Attach the last path. After the loop the stack still holds the final path.
pop_towith size 1 attaches all of it and leaves only the root. - Return the root.
stack.pop()already returns anOption, so it needs no wrapping. - Write
pop_to. It takes the stack as&mut Vec<Box<TreeNode>>and asize: usize. While the stack is longer thansize, it popschild, borrows the new top asparentwithlast_mut(), and moves the child intoparent.leftif that is empty, or intoparent.rightotherwise.
The solution in Rust
impl Solution {
fn pop_to(stack: &mut Vec<Box<TreeNode>>, size: usize) {
while stack.len() > size {
let child = stack.pop().unwrap();
let parent = stack.last_mut().unwrap();
if parent.left.is_none() {
parent.left = Some(child);
} else {
parent.right = Some(child);
}
}
}
fn recover_from_preorder(traversal: String) -> Option<Box<TreeNode>> {
let bytes = traversal.as_bytes();
let mut stack: Vec<Box<TreeNode>> = Vec::new();
let mut i = 0;
while i < bytes.len() {
let mut depth = 0;
while bytes[i] == b'-' {
depth += 1;
i += 1;
}
let mut value = 0;
while i < bytes.len() && bytes[i] != b'-' {
value = value * 10 + (bytes[i] - b'0') as i32;
i += 1;
}
Self::pop_to(&mut stack, depth);
stack.push(Box::new(TreeNode::new(value)));
}
Self::pop_to(&mut stack, 1);
stack.pop()
}
}
Rust details worth noticing:
- Bytes, not characters. Indexing a
Stringby number isn't allowed in Rust, but a&[u8]can be indexed. That is safe here because every character is ASCII. The literalb'-'is au8, so it compares directly withbytes[i]. - Digit arithmetic.
bytes[i] - b'0'is au8between 0 and 9, andas i32converts it to the type ofvalue. Values go up to 1 000 000 000, which fits in ani32. - Inferred types.
depthhas no annotation, but passing it topop_to, which wants ausize, makes it ausize.valuebecomes ani32from the addition. - Passing the stack.
&mut stacklends the vector to the helper for the duration of the call, andstackis usable again afterwards. The helper has noself, so it is called asSelf::pop_to. - Moving a box.
parent.left = Some(child)moveschildinto the parent. After that linechildcan't be used again, which is why theifandelsebranches each use it once. last_mut(). It givesOption<&mut Box<TreeNode>>. Theunwrap()is safe because the loop only runs while the stack is longer thansize, andsizeis at least 1 whenever a pop happens: the root is never popped.
Complexity
Time: O(n), where n is the length of the string. Each character is read once, and each node is pushed once and popped once, with constant work for each.
Space: O(h) beyond the tree itself, where h is the tree's height, because the stack holds one root-to-node path. For a tree that is a single chain, that is the number of nodes.
Edge cases and common mistakes
- A single node.
"7"pushes one node, the finalpop_to(.., 1)does nothing, and the root is returned. - Multi-digit values. Reading one byte as one value breaks
"5-12--7". Keep reading until the next dash or the end. - Two children at the same depth.
"1-2-3"gives the root 1 with 2 on the left and 3 on the right. Popping 2 fills the empty left slot, and popping 3 finds it occupied. - A jump up by several levels. After
---4comes-5: depth goes from 3 to 1. Popping a single node isn't enough, sopop_toloops until the stack has the right size. - Forgetting the final attachment. Without the last
pop_to, the nodes of the last path stay unattached, andstack.pop()returns the last node read instead of the root. - Checking
iin the dash loop. It isn't needed, but in the digit loopi < bytes.len()must come first, or the last value runs past the end of the string and panics.
Testing
"1-2--3--4-5--6--7"gives[1,2,5,3,4,6,7]."1-2--3---4-5--6---7"gives[1,2,5,3,null,6,null,4,null,7]."5-12--7---40--9"gives[5,12,null,7,9,40]."7"gives[7]."1-2-3"gives[1,2,3].
