Problem overview
The function zigzag_traversal receives the root of a binary tree as an Option<Box<TreeNode>> and returns the node values as a Vec<Vec<i32>>, one inner vector per level. The first level is read from left to right, the next from right to left, and the direction keeps alternating. An empty tree (None) gives an empty result.
For the tree with root 3, children 9 and 20, and 20's children 15 and 7, the answer is [[3], [20, 9], [15, 7]]. Level 0 reads left to right, level 1 right to left, and level 2 left to right again.
Approach: a stack reverses the order for free
An ordinary level-order traversal reads each level left to right. Zigzag needs every other level reversed. One way is to collect each level and reverse it afterwards. This lesson avoids the reversal by using the order in which a stack hands things back.
A Vec used as a stack gives back the last element pushed first. If a level's nodes are pushed onto a stack in left-to-right order, they pop right to left. So the trick is to decide, while processing one level, in which order its children go onto the stack for the next level:
- On an even level (read left to right), push the left child and then the right child. The next level is odd and must be read right to left, which is exactly the pop order.
- On an odd level (read right to left), push the right child and then the left child. The next level is even and pops left to right.
Two stacks are needed because a level must be fully read before the next one starts. current_stack holds the level being read, and next_stack collects its children. When current_stack runs empty, the two swap roles.
Step by step
- Create
resultas an empty, mutableVec<Vec<i32>>and return it at once ifrootisNone. The traversal below assumes at least one node. - Declare the two stacks as
Vec<Box<TreeNode>>. They own the nodes they hold, so the element type is the box itself, not a reference. - Seed
current_stackwith the root.if let Some(node) = roottakes the node out of theOptionand moves it into the stack. - Keep a
levelcounter of typeusize, starting at 0. It is used as an index intoresult, and indexing needs ausize. - Loop while
current_stackis not empty. Each pass processes one node. - Pop the top node.
popreturns anOption, andunwrapis safe because the loop condition just checked that the stack has something in it. - Make sure the row for this level exists. The first node of a level finds
level >= result.len()and pushes an empty vector. Later nodes of the same level skip this. - Append the node's value to
result[level]. - Push the children in the order that depends on the level's parity, as described above.
if let Some(l) = node.leftskips a missing child without a separate check. - Detect the end of the level. When
current_stackis empty after handling a node, every node of the level has been read. Swap the stacks withstd::mem::swapand add 1 tolevel. The old stack is empty, so it becomes the collector for the following level. - Return
resultas the final expression.
The solution in Rust
impl Solution {
fn zigzag_traversal(root: Option<Box<TreeNode>>) -> Vec<Vec<i32>> {
let mut result: Vec<Vec<i32>> = vec![];
if root.is_none() {
return result;
}
let mut current_stack: Vec<Box<TreeNode>> = vec![];
let mut next_stack: Vec<Box<TreeNode>> = vec![];
if let Some(node) = root {
current_stack.push(node);
}
let mut level = 0usize;
while !current_stack.is_empty() {
let node = current_stack.pop().unwrap();
if level >= result.len() {
result.push(vec![]);
}
result[level].push(node.val);
if level % 2 == 0 {
if let Some(l) = node.left {
next_stack.push(l);
}
if let Some(r) = node.right {
next_stack.push(r);
}
} else {
if let Some(r) = node.right {
next_stack.push(r);
}
if let Some(l) = node.left {
next_stack.push(l);
}
}
if current_stack.is_empty() {
std::mem::swap(&mut current_stack, &mut next_stack);
level += 1;
}
}
result
}
}
The Rust details worth a closer look:
- Moving children out of the node.
nodeis aBox<TreeNode>that the function owns. Writingif let Some(l) = node.leftmoves the left child out of it, which Rust allows for a box. The box is never used as a whole again, butnode.valwas already copied intoresultearlier, andi32isCopyanyway. Nothing is cloned, so the tree is consumed as it is read. - Why
unwrapis fine here.unwraponNonepanics, but this call followswhile !current_stack.is_empty(), so it always receivesSome. std::mem::swap. It takes two&mutreferences and exchanges the values behind them without copying the nodes: only the vectors' internal pointers change places.- Parity with
%.level % 2 == 0is true for 0, 2, 4 and so on. The counter is ausize, so it is never negative and the remainder is never surprising. - Final expression.
resulthas no semicolon, so it is the return value. The earlyreturn result;does need one.
Complexity
Time: O(n) for a tree with n nodes. Each node is pushed onto a stack once and popped once, and the work per node is constant.
Space: O(n). The result holds all n values. The stacks together hold at most two adjacent levels, which in the worst case is also proportional to n.
Edge cases and common mistakes
- Empty tree.
Nonereturns[]. Without the early return, the code would still push nothing, but the check makes the intent clear. - Single node. Returns
[[1]]. The stacks are swapped once, both are empty, and the loop ends. - Missing children. A node with only a right child, such as a tree 1 → 2 → 3 where each node has just one child, still works:
if letskips the absent side. A chain gives one value per level. - Pushing in the same order on every level. That produces ordinary level order, with every level read in the same direction. The alternation in the push order is the whole idea.
- Swapping too early. The swap must happen only when
current_stackis empty. Swapping after every node would mix two levels. - Forgetting
level += 1. All values would be appended to the same row. - Reading the direction from the wrong parity. The push order is chosen by the level being read, not by the level being built. Even levels push left first, because the next level is the odd one.
Testing
- Root 3 with children 9 and 20, and 20's children 15 and 7:
[[3], [20, 9], [15, 7]]. - A perfect tree 1; 2, 3; 4, 5, 6, 7:
[[1], [3, 2], [4, 5, 6, 7]]. - A single node 1:
[[1]]. - An empty tree:
[]. - A chain 1 → 2 → 3 down the left side:
[[1], [2], [3]].
