LESSON 309 OF 348 · BINARY TREES

Zigzag Tree Traversal in Kotlin: Two Stacks, Level by Level

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 16

Guided

Ctrl + Enter to check

KOTLIN · BINARY TREES

Zigzag Tree Traversal in Kotlin, explained

Problem overview

The function zigzagTraversal takes the root of a binary tree, a TreeNode?, and returns the values level by level as a List<List<Int>>. The root's level is read from left to right, the next from right to left, and so on. 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]]. An empty tree gives an empty list.

Approach: let the stacks reverse each level

The usual way to read a tree level by level is a queue, which hands nodes back in the order they were added. A stack does the opposite: it hands back the most recently added node first. That reversal is what a zigzag needs, if the children are added in the right order.

The solution uses two stacks. currentLevelStack holds the nodes of the level being read, and each node popped from it is the next value of that level. nextLevelStack collects the children of those nodes. When the current stack runs empty, the level is complete, the next stack becomes the current one, and the level number goes up.

The order in which children are pushed decides how the following level is read:

On an even level the children are pushed left first, then right. The right child is on top, so the next level pops right to left, which is the direction an odd level needs.
On an odd level they are pushed right first, then left. The left child is on top, so the next level pops left to right.

On the example, level 0 holds only 3. It pushes 9 and then 20, so level 1 pops 20 first and then 9. Reading 20 pushes its children as right then left, which is 7 then 15; reading 9 pushes nothing. Level 2 therefore pops 15 and then 7, giving [15, 7]. The whole traversal needs no reversing step at all.

The first lesson on this problem does the opposite: it walks the tree recursively, fills the levels left to right and reverses the odd ones afterwards. This lesson builds the right order directly.

Step by step

Create result as an empty mutable list of mutable lists of Int, with one inner list per level.
Create the two stacks as vars holding empty ArrayDeque<TreeNode> objects. They must be var because they are swapped later. An ArrayDeque works as a stack when you only add and remove at its last end.
Put the root on the current stack with a helper, addToStack, that pushes only the nodes that aren't null. For an empty tree nothing is pushed and the loop never runs.
Track the level in a var level that starts at 0.
Loop while the current stack is not empty. Each iteration handles one node.
Pop a node with removeLast().
Record its value with a second helper, addToResult, that appends to the list for the given level and creates that list when the level is new.
On an even level, push the left child and then the right child onto the next stack.
On an odd level, push the right child and then the left child.
Finish the level when the current stack is empty. Make the next stack the current one, give the next stack a new empty deque, and add 1 to level.
Return result after the loop.
Write addToResult. If level >= result.size, add a new empty list. This is always enough, because the first node of each level arrives before any node of the next one. Then add the value to result[level].
Write addToStack. It takes the stack and a list of nullable nodes, and for each node that isn't null it calls addLast. The children are passed in without checking them, so this helper is where missing children are filtered out.

The solution in Kotlin

class Solution {

    fun addToStack(stack: ArrayDeque<TreeNode>, nodes: List<TreeNode?>) {
        for (node in nodes) {
            if (node != null) {
                stack.addLast(node)
            }
        }
    }

    fun addToResult(result: MutableList<MutableList<Int>>, level: Int, value: Int) {
        if (level >= result.size) {
            result.add(mutableListOf())
        }
        result[level].add(value)
    }

    fun zigzagTraversal(root: TreeNode?): List<List<Int>> {
        val result = mutableListOf<MutableList<Int>>()
        var currentLevelStack = ArrayDeque<TreeNode>()
        var nextLevelStack = ArrayDeque<TreeNode>()

        addToStack(currentLevelStack, listOf(root))
        var level = 0

        while (currentLevelStack.isNotEmpty()) {
            val node = currentLevelStack.removeLast()

            addToResult(result, level, node.value)

            if (level % 2 == 0) {
                addToStack(nextLevelStack, listOf(node.left, node.right))
            } else {
                addToStack(nextLevelStack, listOf(node.right, node.left))
            }

            if (currentLevelStack.isEmpty()) {
                currentLevelStack = nextLevelStack
                nextLevelStack = ArrayDeque()
                level++
            }
        }

        return result
    }
}

The Kotlin details worth knowing:

Nullable and non-null element types. The stacks hold TreeNode, not TreeNode?. The helper takes List<TreeNode?> and, inside if (node != null), the compiler treats node as a TreeNode, so addLast(node) is accepted. Because only real nodes get in, removeLast() returns a non-null node and node.value needs no ?..
listOf(root) with a nullable root. It creates a List<TreeNode?>. If root is null, the list holds a single null, and the helper skips it.
Kotlin's ArrayDeque. It is in the standard library and needs no import. addLast and removeLast make it a stack, and removeLast() on an empty deque throws an exception, which the isNotEmpty() condition avoids.
nextLevelStack = ArrayDeque(). The type argument is inferred from the declared type of the variable, so it isn't written again.
Reassigning references. currentLevelStack = nextLevelStack makes both names point at the same object, and the next line gives nextLevelStack a new one. The swap is cheap because no elements are copied.
The remainder operator. level % 2 == 0 is true for even levels. Levels are never negative, so the sign behavior of % doesn't matter.

Complexity

Let n be the number of nodes.

Time: O(n). Every node is pushed once and popped once, and each pop does a constant amount of work: one value added to a list and at most two children pushed.

Space: O(n). The result holds all values. Besides it, the two stacks together hold at most two adjacent levels, which is no more than n nodes and is typically about the width of the tree.

Edge cases and common mistakes

Empty tree. addToStack skips the null root, the loop doesn't run, and the result is [].
Single node. The result is [[1]]. The node is popped, no children are pushed, and the current stack is empty, so the swap happens once and the loop ends.
Pushing children in the wrong order. Left-then-right on every level gives the right-to-left reading on every level. The order must alternate, and the level's parity decides it.
Swapping the stacks too early. The swap is only safe when the current stack is empty. Doing it before the level is finished would mix the two levels.
Forgetting to start a fresh next stack. If both names pointed at the same deque after the swap, the next level's children would be pushed onto the stack that is being read.
Not incrementing level. All values would pile into the first list, and the direction would never alternate.
Missing children. A node with one child needs no special case, because the helper drops the null.

Testing

The tree [3, 9, 20, null, null, 15, 7] in level order gives [[3], [20, 9], [15, 7]].
The tree [1, 2, 3, 4, 5, 6, 7] gives [[1], [3, 2], [4, 5, 6, 7]].
A single node [1] gives [[1]].
An empty tree gives [].

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

Start learning free