Binary Tree BFS: Zigzag Order

Written By James Bowen

In our last article, we explored how to perform an in-order traversal of a binary search tree. Today we’ll do one final binary tree problem to solidify our understanding of some common tree patterns, as well as the tricky syntax for dealing with a binary tree in Rust.

If you want some interesting challenge problems using Haskell data structures, you should take our Solve.hs course. In particular, you’ll learn how to write a self-balancing binary tree to use for an ordered set!

The Problem

Today we will solve Zigzag Level Order Traversal. For any binary tree, we can think about it in terms of “levels” based on the number of steps from the root. So given this tree:

-    45
    /  \
   32  50
  /  \   \
 5   40   100
    /  \
  37   43

We can visually see that there are 4 levels. So a normal level order traversal would return a list of 4 lists, where each list is a single level, ordered from left to right, visually speaking:

[45]
[32, 50]
[5, 40, 100]
[37, 43]

However, with a zigzag level order traversal, every other level is reversed. So we should get the following result for the input tree:

[45]
[50, 32]
[5, 40, 100]
[43, 37]

So we can imagine that we do the first level from left to right and then zigzag back to get the second level from right to level. Then we do left to right again for the third level, and so on.

The Algorithm

For our in-order traversal, we used a kind of depth-first search (DFS), and this approach is more common for tree-based problems. However, for a level-order problem, we want more of a breadth-first search (BFS). In a BFS, we explore states in order of their distance to the root. Since all nodes in a level have the same distance to the root, this makes sense.

Our general idea is that we’ll store a list of all the nodes from the prior level. Initially, this will just contain the root node. We’ll loop through this list, and create a new list of the values from the nodes in this list. This gets appended to our final result list.

While we’re doing this loop, we’ll also compose the list for the next level. The only trick is knowing whether to add each node’s left or right child to the next-level list first. This flips each iteration, so we’ll need a boolean tracking it that flips each time.

Once we encounter a level that produces no numbers (i.e. it only contains Nil nodes), we can stop iterating and return our list of lists.

Rust Solution

Now that we’re a bit more familiar with manipulating Rc RefCells, we’ll start with the Rust solution, framing it according to the two-loop structure in our algorithm. We’ll define stack1, which is the iteration stack, and stack2, where we accumulate the new nodes for the next layer. We also define our final result vector, a list of lists.

pub fn zigzag_level_order(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<Vec<i32>> {
    let mut result: Vec<Vec<i32>> = Vec::new();
    let mut stack1: Vec<Option<Rc<RefCell<TreeNode>>>> = Vec::new();
    stack1.push(root.clone());
    let mut stack2: Vec<Option<Rc<RefCell<TreeNode>>>> = Vec::new();
    let mut leftToRight = true;

    ...
    return result;
}

Our initial loop will continue until stack1 no longer contains any elements. So our basic condition is while(!stack1.is_empty(). However, there’s another important element here.

After we accumulate the new nodes in stack2, we want to flip the meanings of our two stacks. We want our accumulated nodes referred to by stack1, and stack2 to be an empty list to accumulate. We accomplish this in Rust by clearing stack1 at the end of our loop, and then using std::mem::swap to flip their meanings:

pub fn zigzag_level_order(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<Vec<i32>> {
    let mut result: Vec<Vec<i32>> = Vec::new();
    let mut stack1: Vec<Option<Rc<RefCell<TreeNode>>>> = Vec::new();
    stack1.push(root.clone());
    let mut stack2: Vec<Option<Rc<RefCell<TreeNode>>>> = Vec::new();
    let mut leftToRight = true;

    while (!stack1.is_empty()) {
        let mut thisLayer = Vec::new(); // Values from this level
        ...
        leftToRight = !leftToRight;
        stack1.clear();
        mem::swap(&mut stack1, &mut stack2);
    }
    return result;
}

In C++ we could accomplish something like this using std::move, but only because we want stack1 to return to an empty state.

stack2 = std::move(stack1);

Also, observe that we flip our boolean flag at the end of the iteration.

Now let’s get to work on the inner loop. This will actually go through stack1, add values to thisLayer, and accumulates the next layer of nodes for stack2. An interesting finding is that whether we’re going left to right or vice versa, we want to loop through stack2 in reverse. This means we’re treating it like a true stack instead of a vector, first accessing the last node to be added.

A left-to-right pass will add lefts and then rights. This means the right-mode node in the next layer is on “top” of the stack, at the end of the vector. A right-to-left pass will first add the right child for a node before its left. This means the left-most node of the next layer is at the end of the vector.

Let’s frame up this loop, and also add the results of this layer to our final result vector.

pub fn zigzag_level_order(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<Vec<i32>> {
    ...

    while (!stack1.is_empty()) {
        let mut thisLayer = Vec::new()
        for node in stack1.iter().rev() {
            ...
        }

        if (!thisLayer.is_empty()) {
            result.push(thisLayer);
        }
        leftToRight = !leftToRight;
        stack1.clear();
        mem::swap(&mut stack1, &mut stack2);
    }
    return result;
}

Note that we do not add the values array if it is empty. We allow ourselves to accumulate None nodes in our stack. The final layer we encounter will actually consist of all None nodes, and we don’t want this layer to add an empty list.

Now all we need to do is populate the inner loop. We only take action if the node from stack1 is Some instead of None. Then we follow a few simple steps:

  1. Borrow the TreeNode from this RefCell
  2. Push its value onto thisLayer.
  3. Add its children (using clone) to stack2, in the right order.

Here’s the code:

pub fn zigzag_level_order(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<Vec<i32>> {
    ...

    while (!stack1.is_empty()) {
        let mut thisLayer = Vec::new()
        for node in stack1.iter().rev() {
            if let Some(current) = node {
                let currentTreeNode = current.borrow();
                thisLayer.push(currentTreeNode.val);
                if leftToRight {
                    stack2.push(currentTreeNode.left.clone());
                    stack2.push(currentTreeNode.right.clone());
                } else {
                    stack2.push(currentTreeNode.right.clone());
                    stack2.push(currentTreeNode.left.clone());
                } 
            }
        }

        ...
    }
    return result;
}

And now we’re done! Here’s the full solution:

use std::rc::Rc;
use std::cell::RefCell;
use std::mem;

pub fn zigzag_level_order(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<Vec<i32>> {
    let mut result: Vec<Vec<i32>> = Vec::new();
    let mut stack1: Vec<Option<Rc<RefCell<TreeNode>>>> = Vec::new();
    stack1.push(root.clone());
    let mut stack2: Vec<Option<Rc<RefCell<TreeNode>>>> = Vec::new();

    let mut leftToRight = true;
    while (!stack1.is_empty()) {
        let mut thisLayer = Vec::new();
        for node in stack1.iter().rev() {
            if let Some(current) = node {
                let currentTreeNode = current.borrow();
                thisLayer.push(currentTreeNode.val);
                if leftToRight {
                    stack2.push(currentTreeNode.left.clone());
                    stack2.push(currentTreeNode.right.clone());
                } else {
                    stack2.push(currentTreeNode.right.clone());
                    stack2.push(currentTreeNode.left.clone());
                } 
            }
        }
            
        if (!thisLayer.is_empty()) {
            result.push(thisLayer);
        }
        leftToRight = !leftToRight;
        stack1.clear();
        mem::swap(&mut stack1, &mut stack2);
    }
    return result;
}

Haskell Solution

While our Rust solution was better described from the outside in, it’s easy to build the Haskell solution from the inside out. We have two loops, and we can start by defining the inner loop (we’ll call it the stack loop).

The goal of this loop is to take stack1 and turn it into stack2 (the next layer) and the numbers for this layer, while also tracking the direction of iteration. Both outputs are accumulated as lists, so we have inputs for them as well:

zigzagOrderTraversal :: TreeNode -> [[Int]]
zigzagOrderTraversal root = ...
  where
    stackLoop :: Bool -> [TreeNode] -> [TreeNode] -> [Int] -> ([TreeNode], [Int])
    stackLoop isLeftToRight stack1 stack2 nums = ...

When stack1 is empty, we return our result from this loop. Because of list accumulation order, we reverse nums when giving the result. However, we don’t reverse stack2, because we want to iterate starting from the “top”. This seems like the opposite of what we did in Rust, because Rust uses a vector for its stack type, instead of a singly linked list!

zigzagOrderTraversal :: TreeNode -> [[Int]]
zigzagOrderTraversal root = ...
  where
    stackLoop :: Bool -> [TreeNode] -> [TreeNode] -> [Int] -> ([TreeNode], [Int])
    stackLoop _ [] stack2 nums = (stack2, reverse nums)
    stackLoop isLeftToRight (Nil : rest) stack2 numbers = stackLoop isLeftToRight rest stack2 numbers
    stackLoop isLeftToRight (Node x left right : rest) stack2 nums = ...

Observe also a second edge case…for Nil nodes in stack1, we just recurse on the rest of the list. Now for the main case, we just define the new stack2, which adds the child nodes in the correct order. Then we recurse while also adding x to nums.

zigzagOrderTraversal :: TreeNode -> [[Int]]
zigzagOrderTraversal root = ...
  where
    stackLoop :: Bool -> [TreeNode] -> [TreeNode] -> [Int] -> ([TreeNode], [Int])
    stackLoop _ [] stack2 nums = (stack2, reverse nums)
    stackLoop isLeftToRight (Nil : rest) stack2 numbers = stackLoop isLeftToRight rest stack2 numbers
    stackLoop isLeftToRight (Node x left right : rest) stack2 nums =
      let stack2' = if isLeftToRight then right : left : stack2 else left : right : stack2
      in  stackLoop isLeftToRight rest stack2' (x : nums)

    ...

Now we’ll define the outer loop, which we’ll call the layerLoop. This takes the direction flag and stack1, plus the accumulator list for the results. It also has a simple base case to reverse the results list once stack1 is empty.

zigzagOrderTraversal :: TreeNode -> [[Int]]
zigzagOrderTraversal root = layerLoop True [root] []
  where
    stackLoop :: Bool -> [TreeNode] -> [TreeNode] -> [Int] -> ([TreeNode], [Int])
    stackLoop = ...

    layerLoop :: Bool -> [TreeNode] -> [[Int]] -> [[Int]]
    layerLoop _ [] allNums = reverse allNums
    layerLoop isLeftToRight stack1 allNums = ...

Now in the recursive case, we call the stackLoop to get our new numbers and the stack for the next layer (which we now think of as our new stack1). We then recurse, flipping the boolean flags and adding these new numbers to our results, but only if the list is not empty.

zigzagOrderTraversal :: TreeNode -> [[Int]]
zigzagOrderTraversal root = layerLoop True [root] []
  where
    stackLoop :: Bool -> [TreeNode] -> [TreeNode] -> [Int] -> ([TreeNode], [Int])
    stackLoop = ...

    layerLoop :: Bool -> [TreeNode] -> [[Int]] -> [[Int]]
    layerLoop _ [] allNums = reverse allNums
    layerLoop isLeftToRight stack1 allNums =
      let (stack1', newNums) = stackLoop isLeftToRight stack1 [] []
      in  layerLoop (not isLeftToRight) stack1' (if null newNums then allNums else newNums : allNums)

The last step, as you seen is calling layerLoop from the start with root. We’re done! Here’s our final implementation:

zigzagOrderTraversal :: TreeNode -> [[Int]]
zigzagOrderTraversal root = layerLoop True [root] []
  where
    stackLoop :: Bool -> [TreeNode] -> [TreeNode] -> [Int] -> ([TreeNode], [Int])
    stackLoop _ [] stack2 nums = (stack2, reverse nums)
    stackLoop isLeftToRight (Nil : rest) stack2 numbers = stackLoop isLeftToRight rest stack2 numbers
    stackLoop isLeftToRight (Node x left right : rest) stack2 nums =
      let stack2' = if isLeftToRight then right : left : stack2 else left : right : stack2
      in  stackLoop isLeftToRight rest stack2' (x : nums)
    
    layerLoop :: Bool -> [TreeNode] -> [[Int]] -> [[Int]]
    layerLoop _ [] allNums = reverse allNums
    layerLoop isLeftToRight stack1 allNums =
      let (stack1', newNums) = stackLoop isLeftToRight stack1 [] []
      in  layerLoop (not isLeftToRight) stack1' (if null newNums then allNums else newNums : allNums)

Conclusion

That’s all we’ll do for binary trees right now. In the coming articles we’ll continue to explore more data structures as well as some common algorithms. If you want to learn more about data structures in algorithms in Haskell, check out our course Solve.hs. Modules 2 & 3 are filled with this sorts of content, including lots of practice problems.

James Bowen
Next

In-Order Traversal in Haskell and Rust