LESSON 300 OF 338 · BINARY TREES

Zigzag Tree Traversal in Go: Level Order with Two Stacks

Given a binary tree, return the zigzag level order traversal of its nodes' values. (ie, from left to right, then right to left for the next level and alternate between).

Example:

Input: root = [3, 9, 20, null, null, 15, 7]
    3
   / \
  9  20
    /  \
   15   7

Output: 
[
  [3],
  [20,9],
  [15,7]
]

STEP 1 OF 18

Guided

Ctrl + Enter to check

GO · BINARY TREES

Zigzag Tree Traversal in Go, explained

Problem overview

The function zigzagTraversal receives the root of a binary tree, a *TreeNode, and returns its values level by level as a [][]int. The first level is read left to right, the next right to left, and so on, alternating.

For the tree [3, 9, 20, null, null, 15, 7], node 3 is the root, 9 and 20 are its children, and 20 has children 15 and 7. Level 0 is [3], level 1 read right to left is [20, 9], and level 2 read left to right is [15, 7]. The answer is [[3], [20, 9], [15, 7]]. An empty tree gives an empty result.

Approach: two stacks, one per level

A normal level-order traversal uses a queue. Zigzag order needs each other level reversed, and a stack gives that reversal for free: whatever is pushed last is popped first.

The solution keeps two stacks. currentLevelStack holds the nodes of the level being read, and nextLevelStack collects their children. Popping a node from the current stack gives it in the right reading order for this level. Its children are pushed onto the next stack in an order chosen so that they pop in the opposite direction:

On an even level, read left to right, the next level must pop right to left. Push the left child first, then the right one, so the right child ends up on top.
On an odd level, read right to left, the next level must pop left to right. Push the right child first, then the left one.

When the current stack is empty, the level is done. The next stack becomes the current one and the level counter goes up.

On the example: level 0 pops 3 and pushes 9, then 20, so 20 is on top. Level 1 pops 20 first, recording it, and pushes its children 7 then 15 (odd level, right first). It then pops 9, which has no children. Level 2 starts with 15 on top, so it records 15 and then 7.

The alternative, a queue plus reversing the odd rows, does the same work in a different way. This lesson avoids the extra reversal by choosing the push order.

Step by step

Create result as an empty [][]int, one inner slice per level.
Create both stacks as empty []*TreeNode slices. A slice is a stack when you push with append and pop from the end.
Push the root through the helper addToStack, which skips nil nodes. For an empty tree nothing is pushed, the loop never runs, and the empty result comes back.
Start level at 0.
Loop while len(currentLevelStack) > 0. Each iteration handles one node.
Pop the top node: read the last element, then reslice without it.
Record the value with the helper addToResult(result, level, node.Val), which returns the updated result.
Even level: push node.Left, then node.Right onto the next stack.
Odd level: push node.Right, then node.Left.
End of level. If the current stack is now empty, set it to the next stack, give the next stack a fresh empty slice, and increase level.
Return result after the loop.
Write addToResult. It takes the rows, a level and a value, and returns the rows.
Create the row if needed. If level >= len(result), append an empty []int{}. Nodes arrive level by level, so only the first node of a level triggers this.
Append the value to result[level].
Return result.
Write addToStack. It takes a stack and a variable number of nodes and returns the stack.
Skip missing nodes. Loop over nodes and append each one that is not nil.
Return the stack.

The solution in Go

func zigzagTraversal(root *TreeNode) [][]int {
	result := [][]int{}
	currentLevelStack := []*TreeNode{}
	nextLevelStack := []*TreeNode{}
	currentLevelStack = addToStack(currentLevelStack, root)
	level := 0
	for len(currentLevelStack) > 0 {
		node := currentLevelStack[len(currentLevelStack)-1]
		currentLevelStack = currentLevelStack[:len(currentLevelStack)-1]
		result = addToResult(result, level, node.Val)
		if level%2 == 0 {
			nextLevelStack = addToStack(nextLevelStack, node.Left, node.Right)
		} else {
			nextLevelStack = addToStack(nextLevelStack, node.Right, node.Left)
		}
		if len(currentLevelStack) == 0 {
			currentLevelStack = nextLevelStack
			nextLevelStack = []*TreeNode{}
			level++
		}
	}
	return result
}

func addToResult(result [][]int, level int, value int) [][]int {
	if level >= len(result) {
		result = append(result, []int{})
	}
	result[level] = append(result[level], value)
	return result
}

func addToStack(stack []*TreeNode, nodes ...*TreeNode) []*TreeNode {
	for _, node := range nodes {
		if node != nil {
			stack = append(stack, node)
		}
	}
	return stack
}

The Go details worth a closer look:

Helpers return the slice. A function receives a copy of the slice header. After append adds an element, the header (length, possibly a new array) is different, and the change is visible to the caller only through a returned value. That is why addToResult and addToStack return the slice and the callers assign it back.
Variadic parameters. In nodes ...*TreeNode, the ... lets the caller pass any number of nodes. Inside the function, nodes is a []*TreeNode. The calls addToStack(currentLevelStack, root) and addToStack(nextLevelStack, node.Left, node.Right) pass one and two.
Nil children are passed on purpose. A missing child is a nil pointer. Passing it in and filtering inside the helper keeps the loop free of if node.Left != nil checks. A nil root is handled by the same filter.
Why a fresh slice for the next level. After currentLevelStack = nextLevelStack, both names hold the same nodes. Without a new empty slice, nextLevelStack would keep the finished level's nodes, and the next level's children would be appended after them. When that stack became the current one, the old nodes would be read again.
A slice of slices. result[level] = append(result[level], value) grows one row and stores it back. Each row is its own slice.
level%2. The remainder is 0 for even levels and 1 for odd ones, and the level is never negative here.

Complexity

Time: O(n). Each node is pushed once and popped once, and each push, pop and row append takes constant amortized time.

Space: O(n). The output holds all n values. The stacks hold at most one level each, which can be about n / 2 nodes in a wide tree.

Edge cases and common mistakes

Empty tree. The root is nil, addToStack pushes nothing, and the function returns []. Appending the root directly would put a nil on the stack, and node.Val would panic.
Wrong push order. The children are pushed in the order that makes the next level come out reversed, not the current one. Swapping the two branches gives the traversal that starts right to left.
Single child. Missing nodes are filtered out, so a level may have fewer nodes than twice the previous one. The tree [1,2,null,3] gives [[1], [2], [3]].
Reusing the stack. Skipping the line that gives nextLevelStack a new slice leaves the finished level's nodes in it, so they are read again and the traversal never ends.
Forgetting to advance the level. Without level++, all values end up in the first row.
Perfect tree. For [1,2,3,4,5,6,7], the answer is [[1], [3, 2], [4, 5, 6, 7]]: the last row reads left to right because level 2 is even.

Testing

[3, 9, 20, null, null, 15, 7] gives [[3], [20, 9], [15, 7]].
[1] gives [[1]].
An empty tree gives [].
[1,2,3,4,5,6,7] gives [[1], [3, 2], [4, 5, 6, 7]].

Keep going with GoAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free