Start learning free

LESSON 333 OF 415 · BINARY TREES

Post Order Tree Traversal in Swift: Iterative with an Explicit Stack

Practice O(n) Read the full explanation

Given a binary tree, return the postorder traversal of its nodes' values.

Example:

Input: [1,nil,2,3]
   1
    \
     2
    /
   3

Output: [3,2,1]

Follow up: Recursive solution is trivial, could you do it iteratively?

STEP 1 OF 14
    Ctrl + Enter to check

    SWIFT · BINARY TREES

    Post Order Tree Traversal in Swift, explained

    Problem overview

    The method postOrderTraversal receives the root of a binary tree as a TreeNode? and returns the node values in post order: left subtree, right subtree, then the node. In the example, the root 1 has only a right child, 2, whose left child is 3, and the answer is [3, 2, 1].

    The recursive solution is short. The follow-up in the statement asks for an iterative one, with a loop and no recursive calls. That is this lesson.

    Approach: simulate the call stack yourself

    A recursive function leaves a trail: each active call remembers which node it is on and what it will do next. The three actions of a post order call are, in order: go left, go right, record the node. The computer's call stack stores that trail. To go iterative, build the same trail by hand.

    The stack holds State objects. Each one stands for one call and has two fields: the node and a step, the action the call performs next:

    • 0, LEFT_CALL: start the call for the left child;
    • 1, RIGHT_CALL: start the call for the right child;
    • 2, APPEND_TO_TRAVERSAL: add the node's value to the answer;
    • anything above 2: the call is over.

    Each loop turn takes the top entry, pops it, and does one action. If the call isn't finished, the same node goes back on the stack with its step increased by one. Then, for the left and right actions, a new entry for the child is pushed on top of it. The child's entry is therefore handled first, completely, and the parent resumes only when the child's entries are all popped. That is how recursion returns to its caller.

    Here is the tree with root 1 and left child 2, written as (node, step) entries. Starting from [(1,0)], handling (1,0) leaves [(1,1), (2,0)]. Handling (2,0) leaves [(1,1), (2,1), (nil,0)]. The nil entry is popped and skipped. Then (2,1) becomes (2,2) plus a skipped nil for the right child, and (2,2) appends 2. The same happens for node 1 and appends 1. The answer is [2, 1].

    Step by step

    1. Define the class State inside Solution, with var node: TreeNode?, var step: Int and an initializer init(_ node: TreeNode?, _ step: Int). The node is optional because entries are also pushed for missing children.
    2. Name the steps with three constants: LEFT_CALL is 0, RIGHT_CALL is 1, APPEND_TO_TRAVERSAL is 2. Names are easier to read than bare numbers.
    3. Create traversal, an empty [Int], for the output.
    4. Create stack, an empty [State]. The end of the array is the top.
    5. Push the first call: State(root, LEFT_CALL).
    6. Loop while the stack isn't empty.
    7. Read the top entry into node and step.
    8. Pop it. If the call isn't finished it will be pushed back in a moment, with its next step.
    9. Skip finished entries. A nil node is an empty subtree, and a step above 2 means all three actions were done. Either way, continue.
    10. Push the same node with step + 1. It must go on the stack before any child entry, so that the child ends up on top.
    11. If the step is LEFT_CALL, push the entry for node!.left, starting at LEFT_CALL.
    12. If it is RIGHT_CALL, push the entry for node!.right, also starting at LEFT_CALL, because every call begins by going left.
    13. If it is APPEND_TO_TRAVERSAL, append node!.val to traversal. Both subtrees are finished by now, which is post order.
    14. Return traversal once the stack is empty.

    The solution in Swift

    class Solution {
    class State {
        var node: TreeNode?
        var step: Int
        init(_ node: TreeNode?, _ step: Int) {
            self.node = node
            self.step = step
        }
    }

    func postOrderTraversal(_ root: TreeNode?) -> [Int] {
        let LEFT_CALL = 0
        let RIGHT_CALL = 1
        let APPEND_TO_TRAVERSAL = 2
        var traversal = [Int]()
        var stack = [State]()
        stack.append(State(root, LEFT_CALL))

        while !stack.isEmpty {
            let node = stack.last!.node
            let step = stack.last!.step
            stack.removeLast()

            if node == nil || step > APPEND_TO_TRAVERSAL {
                continue
            }

            stack.append(State(node, step + 1))

            if step == LEFT_CALL {
                stack.append(State(node!.left, LEFT_CALL))
            } else if step == RIGHT_CALL {
                stack.append(State(node!.right, LEFT_CALL))
            } else if step == APPEND_TO_TRAVERSAL {
                traversal.append(node!.val)
            }
        }

        return traversal
    }
}

    Swift details worth a closer look:

    • A nested class. State is declared inside Solution, so code in Solution uses it as State. Its stored properties have no default values, so the initializer must set both. Inside it, self.node = node assigns the parameter to the property of the same name.
    • Class versus struct. Classes are reference types. Here no entry is changed after it is created (a new one is pushed instead), so a struct would behave the same.
    • stack.last!. last returns an optional because an array can be empty. The loop condition guarantees at least one element, so the force-unwrap is safe. The entry is read before removeLast(), which discards the returned value.
    • Optional comparison. node == nil is allowed for TreeNode?, and node! unwraps the value once the nil case has been skipped.
    • Branching on the step. The three actions are chosen with an if / else if chain on integer constants. A switch would work too, as long as it covers every other value with a default case.

    Complexity

    Time: O(n). Each node produces four entries in total (steps 0 to 3), and each missing child produces one entry that is skipped. Every entry is pushed once and popped once, and each pop does constant work.

    Space: O(h) for the stack, plus O(n) for the answer. The stack holds the entries along the current path from the root, so it is as deep as the tree is tall, plus a constant number of extra entries. The result holds all n values.

    Unlike the recursive version, the depth is limited by memory on the heap, not by the call stack. A tree that is a chain of 200,000 nodes is processed without trouble.

    Edge cases and common mistakes

    • An empty tree. The first entry has a nil node, it is skipped, the stack becomes empty and the method returns [].
    • Pushing the parent after the child. If the incremented parent entry were pushed after the child's entry, the parent would be processed first and the order would be wrong. The parent goes first, so the child lies on top.
    • Starting a child at the wrong step. Every child call must begin at LEFT_CALL. Starting it at RIGHT_CALL would skip its left subtree.
    • Forgetting the finished-entry check. The entry with step 3 is pushed after the append action. Without step > APPEND_TO_TRAVERSAL in the skip condition, it would be processed again.
    • Unwrapping a nil node. Using node! before the skip check would crash on an empty subtree.
    • Reaching for a simpler trick. Reversing a pre order with right-before-left also gives post order, but this lesson teaches the general method, which works for other orders by changing which action comes when.

    Testing

    • The tree [1, null, 2, 3] gives [3, 2, 1].
    • An empty tree gives [].
    • The single node [1] gives [1].
    • The tree [1, 2, 3, 4, 5] gives [4, 5, 2, 3, 1].
    Keep going with SwiftAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free