LESSON 334 OF 415 · BINARY TREES

Zigzag Tree Traversal in Swift: 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, nil, nil, 15, 7]
    3
   / \
  9  20
    /  \
   15   7

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

STEP 1 OF 19

Guided

Ctrl + Enter to check

SWIFT · BINARY TREES

Zigzag Tree Traversal in Swift, explained

Problem overview

The method zigzagTraversal receives the root of a binary tree, as a TreeNode?, and returns the node values grouped by level as a [[Int]]. The first level is read from left to right, the second from right to left, the third from left to right again, and so on.

In the example, 3 is the root, its children are 9 and 20, and 20 has the children 15 and 7. Level 0 is [3]. Level 1 is read right to left, so it is [20, 9]. Level 2 goes left to right again: [15, 7]. The answer is [[3], [20, 9], [15, 7]]. An empty tree returns [].

Approach: a stack reverses what you put into it

Reading a tree level by level is the easy part. The zigzag is the twist: every second level comes out reversed. One way is to collect each level normally and reverse the odd ones afterwards. This lesson avoids the reversal by using the order in which a stack gives items back.

A stack returns the item pushed last. If the children of a level are pushed left to right, they pop right to left, and the other way round. So one stack holds the level being read, and a second stack collects the children of the nodes as they are popped. The order in which children are pushed decides the direction of the next level:

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

Because the children go to a different stack, the current level is not disturbed. Only when the current stack runs empty, the whole level has been read, and the second stack becomes the new current stack.

Take the example. Level 0 holds only 3. Popping it, the level is even, so 9 is pushed and then 20; the next stack is [9, 20]. The current stack is empty, so it becomes the current one and the level number goes to 1. Popping gives 20 first, then 9, which is the right-to-left order. While reading 20 on an odd level, its children are pushed right first: 7, then 15. Reading 9 adds nothing. The next stack is [7, 15], and popping it gives 15 then 7.

Step by step

Declare result with var as an empty [[Int]]. It collects one list per level, already in zigzag order.
Declare the two stacks, currentLevelStack and nextLevelStack, as empty arrays of TreeNode. An array is a stack in Swift: append pushes and removeLast pops.
Push the root through a helper, addToStack, called with &currentLevelStack and [root]. The helper only pushes nodes that exist, so a nil root leaves the stack empty and the loop never runs.
Start a counter level at 0.
Loop while the current stack is not empty. Each iteration handles one node.
Pop the top node into a constant node with removeLast().
Record its value with the helper addToResult, passing &result, level and node.val.
On an even level, push the children left first: [node.left, node.right] onto the next stack. The array may contain nil; the helper skips those.
Otherwise push them right first: [node.right, node.left].
When the current stack is empty, switch levels. Assign the next stack to the current one, give the next stack a fresh empty array, and add 1 to level.
Return result after the loop.
Write addToResult. Its first parameter is inout [[Int]], so changes reach the caller.
Create the level's list when it is missing: if level >= result.count, append an empty [Int](). Levels are visited in order, so this happens at the first node of each level.
Append the value to result[level].
Write addToStack with an inout [TreeNode] stack and a [TreeNode?] array of candidates, since missing children are passed in.
Push only real nodes. Loop over the candidates and use if let node = node to unwrap each one, then append it.

The solution in Swift

class Solution {

    func addToStack(_ stack: inout [TreeNode], _ nodes: [TreeNode?]) {
        for node in nodes {
            if let node = node {
                stack.append(node)
            }
        }
    }

    func addToResult(_ result: inout [[Int]], _ level: Int, _ value: Int) {
        if level >= result.count {
            result.append([Int]())
        }
        result[level].append(value)
    }

    func zigzagTraversal(_ root: TreeNode?) -> [[Int]] {
        var result = [[Int]]()
        var currentLevelStack = [TreeNode]()
        var nextLevelStack = [TreeNode]()

        addToStack(&currentLevelStack, [root])
        var level = 0

        while !currentLevelStack.isEmpty {
            let node = currentLevelStack.removeLast()

            addToResult(&result, level, node.val)

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

            if currentLevelStack.isEmpty {
                currentLevelStack = nextLevelStack
                nextLevelStack = [TreeNode]()
                level += 1
            }
        }

        return result
    }
}

Swift details worth knowing:

[TreeNode?] against [TreeNode]. A node's left and right are optional, so the literal [node.left, node.right] is an array of optionals. The stacks hold plain TreeNode values, and unwrapping happens in one place, addToStack.
if let node = node. This shadows the loop variable with a non-optional of the same name, valid only inside the if.
Arrays are values. currentLevelStack = nextLevelStack copies the array (cheaply, Swift copies the storage only when one of them is later modified). Giving nextLevelStack a new empty array afterwards is therefore safe.
inout parameters. Both helpers change the caller's array, so the parameters are inout and the calls pass &. The _ labels let callers leave the argument names out.
removeLast() on an empty array crashes. The loop condition makes sure the stack has an element every time it is called.

Complexity

Time: O(n), where n is the number of nodes. Every node is pushed once and popped once, and each does constant work.

Space: O(n). The result holds every value. The two stacks together never hold more than two adjacent levels, which can reach about n / 2 nodes in a wide tree. There is no recursion, so a very deep tree cannot overflow the call stack.

Edge cases and common mistakes

Empty tree. nil never reaches the stack, so the result is [].
Single node. The tree [1] gives [[1]].
Missing children. A node with only one child works because nil is skipped. The if let in addToStack is what keeps the stacks free of missing nodes.
Pushing in the wrong order. Swapping the two addToStack calls flips the direction of every level, giving [[3], [9, 20], [7, 15]] for the example.
Switching levels too early. The swap must wait until the current stack is empty. Doing it after every node would mix two levels.
Reusing the same stack. Pushing children onto the stack being popped would interleave levels. This is why there are two.

Testing

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

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

Start learning free