1
\$\begingroup\$

I wanted to make a Sequence that can do a pre-order traversal of Binary Tree. Doing so provides automatically unlocks wonderfully useful methods, like prefix, map, filter, etc. But doing so requires implementing an Iterator, which conforms to IteratorProtocol.

Implementing next() to vend elements means that the usual simple recursive solutions aren't possible (because Swift doesn't have continuations, you have no way of "suspending" your recursion until the next call to next()). This is my best attempt. It's working and gives the correct answers, but I would love some feedback.

Here's my primary Binary Tree data structure (note the absence of parent references):

//=== BinaryTreeNode.swift

public final class BinaryTreeNode<T>: BinaryTreeNodeProtocol {
    public let payload: T
    public let left: BinaryTreeNode<T>?
    public let right: BinaryTreeNode<T>?

    init(
        _ payload: T,
        _ left: BinaryTreeNode<T>? = nil,
        _ right: BinaryTreeNode<T>? = nil
        ) {
        self.payload = payload
        self.left = left
        self.right = right
    }
}

To allow this iterator to be used in more cases, I extracted out the binary-freeness into a protocol. For example, UIViews could be made to conform, without requiring people to reproduce portions of their view hierarchy into my BinaryTreeNode objects.

//=== BinaryTreeNodeProtocol.swift

public protocol BinaryTreeNodeProtocol {
    associatedtype Payload
    var payload: Payload { get }
    var left: Self? { get }
    var right: Self? { get }
}

Here's the Sequence/Iterator implementation:

//=== BinaryPreorderIterator.swift

public extension BinaryTreeNodeProtocol {
    var preorderTraversal: BinaryTreePreorderSequence<Self> { return BinaryTreePreorderSequence(self) }
}

public struct BinaryTreePreorderSequence<Tree: BinaryTreeNodeProtocol>: Sequence {
    private let tree: Tree

    public init(_ tree: Tree) { self.tree = tree }

    public func makeIterator() -> BinaryPreorderIterator<Tree> {
        return BinaryPreorderIterator(tree: self.tree)
    }
}

public struct BinaryPreorderIterator<Tree: BinaryTreeNodeProtocol>: IteratorProtocol {
    private var backtrackStack: [(node: Tree, nextAction: NextAction)]

    private enum NextAction { case visitRoot, visitLeft, visitRight, returnToParent }

    public init(tree: Tree) {
        backtrackStack = [(tree, .visitRoot)]
    }

    public mutating func next() -> Tree.Payload? {
        // Using a non-constant boolean expression circumvent's the compilers warnings about the value always being false.
        let debugPrinting = Bool("false")!

        if debugPrinting { print("next()") }

        while let popped = backtrackStack.last {

            if debugPrinting {
                print("\tloop - backtrackStack:")
                for (node, nextAction) in backtrackStack.dropLast() {
                    print("\t\t| \(node.payload) \(nextAction)")
                }
                print("\t\t↳ \(popped.node.payload) \(popped.nextAction) <--- current state\n")
            }

            advanceState()

            switch (popped.nextAction) {
            case .visitRoot:
                let returnValue = popped.node.payload
                if debugPrinting { print("\treturning \(returnValue)\n") }
                return returnValue

            case .visitLeft:
                // This block could be more nicely expressed as the following code,
                // but that doesn't work, because the second method call is unconditional,
                // unlike an if/else-if pair.
                //
                // tryMarkingLeftToVisitNext()
                // tryMarkingRightToVisitNext()

                if let leftChild = lastState().node.left {
                    self.visitNext(leftChild)
                }
                else if let rightChild = lastState().node.right {
                    self.visitNext(rightChild)
                }

            case .visitRight:
                tryMarkingRightToVisitNext()

            case .returnToParent:
                break
            }
        }

        return nil
    }

    private func assertLastStateExists() {
        precondition(!backtrackStack.isEmpty, """
        lastState() should only be called when there IS a last state, i.e. the backtrace stack isn't empty.
        """)
    }

    private func lastState() -> (node: Tree, nextAction: NextAction) {
        assertLastStateExists()
        return backtrackStack.last!
    }

    private mutating func setNextActionOnLastNode(_ newNextAction: NextAction) {
        assertLastStateExists()
        backtrackStack[backtrackStack.count - 1].nextAction = newNextAction // gross, but Array.last is write-only
    }

    private mutating func advanceState() {
        switch lastState().nextAction {
        case .visitRoot:  setNextActionOnLastNode(.visitLeft)
        case .visitLeft:  setNextActionOnLastNode(.visitRight)
        case .visitRight: setNextActionOnLastNode(.returnToParent)
        case .returnToParent: backtrackStack.removeLast()
        }
    }

    private mutating func tryMarkingLeftToVisitNext() {
        if let leftChild = lastState().node.left {
            self.visitNext(leftChild)
        }
    }

    private mutating func tryMarkingRightToVisitNext() {
        if let rightChild = lastState().node.right {
            self.visitNext(rightChild)
        }
    }

    private mutating func visitNext(_ node: Tree) {
        backtrackStack.append((node: node, nextAction: .visitRoot))
    }
}

And here's a little demo:

func sampleTree() -> BinaryTreeNode<Int> {
    typealias B = BinaryTreeNode

    //      4
    //    /   \
    //   2     6
    //  / \   / \
    // 1   3 5   7
    return B(4, B(2, B(1), B(3)), B(6, B(5), B(7)))
}

print(Array(sampleTree().preorderTraversal))

Questions:

  1. Is there a way to extract the if/else if pair into calls to tryMarkingLeftToVisitNext()/tryMarkingRightToVisitNext()? Obviously they could be wrapped in the same if/else if pair, but that doesn't really tidy anything.
  2. I have a strong suspition that one of the NextAction states can be removed entirely.

    For example, if elements' payloads were returned upon visiting them, they could be added to the backtrace with a next state of visitLeft, omitting the need for visitRoot (merely existing in the backtrace implies visitRoot was already done).

    I was close to getting this going, but faced issues with how to visit the root (since my backtrace is initialized to [(tree, .visitRoot)], to kick things off on the first call to next()). I could only achieve it if I added a separate root instance variable to my iterator, and a special case in next() branch to handle the root.

All other feedback and suggestions are greatly appreciated!

\$\endgroup\$
2
\$\begingroup\$

Let's start with a bug: If a node has a right but no left child then the right subtree is traversed twice. As an example, the tree

//      1
//        \
//         2
//        / \
//       3   4
return B(1, nil, B(2, B(3), B(4)))

produces the output

[1, 2, 3, 4, 2, 3, 4]

The reason is that for the (empty) left child of node ①, the right child node ② is added twice to the stack, in func advanceState() at

case .visitLeft:  setNextActionOnLastNode(.visitRight)

and in the “main loop” at

else if let rightChild = lastState().node.right {
    self.visitNext(rightChild)
}

The solution is simple: Remove that else case in case .visitLeft:. Which means that you call just tryMarkingLeftToVisitNext(), as intended:

        case .visitLeft:
            tryMarkingLeftToVisitNext()

        case .visitRight:
            tryMarkingRightToVisitNext()

or inline those two functions, since each of them is called only once.


What I don't like is how the stack is treated in the main loop:

while let popped = backtrackStack.last {
    advanceState()
    switch (popped.nextAction) {
        // ...
    }
}

The variable name popped suggests that the element is removed from the stack, but it isn't: It remains on the stack and is mutated or removed in func advanceState().

I would also suggest to have a single place where the current state is acted upon, not two places (which can cause subtle bugs, as we saw above).

The main loop then would look like this

while let popped = backtrackStack.popLast() {
    switch (popped.nextAction) {
        // ...
    }
}

where the switch handles the different cases, and appends new states to the stack where necessary. This

  • simplifies the code and the logic,
  • makes various utility functions obsolete,
  • makes the .removeToParent case obsolete,
  • makes the forced unwrapping in last! and the checks for a non-empty stack obsolete at various places.

This leads to the following implementation of struct BinaryPreorderIterator:

public struct BinaryPreorderIterator<Tree: BinaryTreeNodeProtocol>: IteratorProtocol {
    private var backtrackStack: [(node: Tree, nextAction: NextAction)]
    private enum NextAction { case visitRoot, visitLeft, visitRight }

    public init(tree: Tree) {
        backtrackStack = [(tree, .visitRoot)]
    }

    public mutating func next() -> Tree.Payload? {
        while let popped = backtrackStack.popLast() {
            switch (popped.nextAction) {
            case .visitRoot:
                backtrackStack.append((popped.node, .visitLeft))
                return popped.node.payload

            case .visitLeft:
                backtrackStack.append((popped.node, .visitRight))
                if let leftChild = popped.node.left {
                    backtrackStack.append((leftChild, .visitRoot))
                }

            case .visitRight:
                if let rightChild = popped.node.right {
                    backtrackStack.append((rightChild, .visitRoot))
                }
            }
        }

        return nil
    }
}

And now we see that the stack actually just identifies the nodes which have to be visited later. If we push left and right child nodes onto the stack immediately when a node is encountered then the enum NextAction and the while-loop are not needed anymore:

public struct BinaryPreorderIterator<Tree: BinaryTreeNodeProtocol>: IteratorProtocol {
    private var backtrackStack: [Tree]

    public init(tree: Tree) {
        backtrackStack = [tree]
    }

    public mutating func next() -> Tree.Payload? {
        guard let node = backtrackStack.popLast() else {
            return nil // Stack is empty.
        }
        if let rightChild = node.right {
            backtrackStack.append(rightChild)
        }
        if let leftChild = node.left {
            backtrackStack.append(leftChild)
        }
        return node.payload
    }
}

Note that we have reduced the implementation of struct BinaryPreorderIterator from 75 lines to 20 lines (not counting the debug output) and simplified it considerably. It also takes less memory and is more efficient.

Remark: func next()could be more compactly written using Optional.map:

    public mutating func next() -> Tree.Payload? {
        guard let node = backtrackStack.popLast() else {
            return nil
        }
        node.left.map { backtrackStack.append($0) }
        node.right.map { backtrackStack.append($0) }
        return node.payload
    }

or even

    public mutating func next() -> Tree.Payload? {
        return backtrackStack.popLast().map { node in
            node.left.map { backtrackStack.append($0) }
            node.right.map { backtrackStack.append($0) }
            return node.payload
        }
    }

but that does not make it better readable or understandable.


Some further suggestions:

It suffices to implement the IteratorProtocol protocol and to declare the Sequence conformance, making struct BinaryTreePreorderSequence obsolete:

public extension BinaryTreeNodeProtocol {
    var preorderTraversal: BinaryPreorderIterator<Self> {
        return BinaryPreorderIterator(tree: self) }
}

public struct BinaryPreorderIterator<Tree: BinaryTreeNodeProtocol>: IteratorProtocol, Sequence {
    private var backtrackStack: [Tree]

    public init(tree: Tree) {
        backtrackStack = [tree]
    }

    public mutating func next() -> Tree.Payload? {
        guard let node = backtrackStack.popLast() else {
            return nil
        }
        if let rightChild = node.right {
            backtrackStack.append(rightChild)
        }
        if let leftChild = node.left {
            backtrackStack.append(leftChild)
        }
        return node.payload
    }
}

This works because there is a default implementation of func makeIterator() for IteratorProtocol.

Finally, I would not use empty argument labels in the init method of class BinaryTreeNode:

public final class BinaryTreeNode<T>: BinaryTreeNodeProtocol {
    // ...

    init(
        _ payload: T,
        left: BinaryTreeNode<T>? = nil,
        right: BinaryTreeNode<T>? = nil
        ) {
        // ...
    }
}

This requires more typing, but may be easier to read, and allows to create nodes with a right subtree only, without having to use a nil argument for the left subtree:

B(1, right: B(2, left: B(3), right: B(4)))
// versus
B(1, nil, B(2, B(3), B(4)))
\$\endgroup\$
  • \$\begingroup\$ Wow, this is much more improved than I was expecting! I wasn't thrilled about the while let popped = backtrackStack.last part either, but I found myself re-appending the elements I just popped off, so I changed to this "peek" approach instead but forgot to rename the popped variable. (Boy I wish we just had a simple stack with nice peek()/push()/pop() operations) \$\endgroup\$ – Alexander May 31 at 14:04
  • \$\begingroup\$ I understand the removal of if let rightChild = popped.node.right under the .visitLeft case (since its absence just caused another loop around, this time reaching the .visitRight state because of the advance() call from the previous loop). I'm a little confused about how you were able to remove the need for the loop, but part of the problem there was that I don't remember the problem that the loop initially solved lol. I'll think on it! \$\endgroup\$ – Alexander May 31 at 14:08
  • \$\begingroup\$ I have mixed feelings about conforming to both IteratorProtocol and Sequence. I knew it was possible, but I thought you would need to add a makeIterator() implementation that just does return self. Didn't know that the standard library adds that by default, so that serves as a kind of "endorsement". But I still don't like the idea of making one-shot sequences where it's not strictly necessary, even if making a separate Sequence requires this extra boilerplate \$\endgroup\$ – Alexander May 31 at 14:13
  • \$\begingroup\$ Think you so much for your time Martin, I appreciate it. I'll work over this and let you know if I have any follow up questions :) \$\endgroup\$ – Alexander May 31 at 14:20
  • \$\begingroup\$ Suppose the binary tree nodes had parent references, it wouldn't be possible to do any better, space-wise, right? Intuitively that makes sense to me, but maybe I'm missing something. \$\endgroup\$ – Alexander Jun 3 at 16:11

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.