Traversing a Binary Search Tree in Swift

Question

I'm implementing the Binary Search Tree data structure in Swift. It looks like this:

class BinarySearchTree<Key: Comparable, Value> {

    let key: Key
    var value: Value
    var left, right: BinarySearchTree<Key, Value>?

    init(key: Key, value: Value) {
        self.key = key
        self.value = value
    }

    // irrelevant methods for constructing a tree

}

To be able to traverse through it using for (key, value) in myTree { }, BinarySearchTree has to implement SequenceType. This is my first attempt:

​
extension BinarySearchTree: SequenceType {

    func generate() -> AnyGenerator<(Key, Value)> {
        let leftGenerator = left?.generate()
        let rightGenerator = right?.generate()

        let (key, value) = (self.key, self.value)
        var generatedSelf = false

        return AnyGenerator {
            if generatedSelf { return rightGenerator?.next() }
            else if let next = leftGenerator?.next() { return next }
            else {
                generatedSelf = true
                return (key, value)
            }
        }
    }

}

It works fine, but when I run benchmarks, it's just not as fast as I'd like. Can my traversal algorithm be improved, or is this the fastest way to do it?

Answer 1

Your generate() method works correctly, as far as I can see. But it creates an AnyGenerator for every node in the tree, i.e. it uses more memory than necessary. Each call to the next() method of the generator then traverses down the next() methods along the left or right subtrees, which makes it slow.

An iterative approach using a stack instead of "nested generators" seems to be faster. Here is a possible implementation of the iterativeInorder algorithm described in Wikipedia: Tree traversal:

func generate() -> AnyGenerator<(Key, Value)> {

    var node : BinarySearchTree? = self
    var stack : [BinarySearchTree] = []

    return AnyGenerator {
        // Traverse from current node to the far left, pushing all 
        // nodes onto the stack:
        while let node1 = node {
            stack.append(node1)
            node = node1.left
        }
        if !stack.isEmpty {
            // Pop node from stack and return its (key, value),
            // then continue with right subtree.
            let current = stack.removeLast()
            defer { node = current.right }
            return (current.key, current.value)
        } else {
            return nil
        }
    }
}

Note that the size of the stack array is limited by the tree height.

In my test with a tree of 1,000,000 nodes, the traversal is done in about 0.2 seconds, compared to 3.5 seconds with your original method (test done in Release mode on a MacBook).