2
\$\begingroup\$

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?

\$\endgroup\$

1 Answer 1

2
\$\begingroup\$

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).

\$\endgroup\$
7
  • \$\begingroup\$ Thanks, your implementation is indeed a lot faster than mine. I also did some benchmarks with a tree of 1,000,000 nodes. My original method takes a little under 4 seconds, similar to your results, but traversal using your method still takes nearly a second for me, and I'm not sure why. I populate the tree with a million random values using arc4random_uniform(1_000_000) (though of course I didn't start measuring until after the tree was populated) - was your method any different? \$\endgroup\$ May 7, 2016 at 12:13
  • \$\begingroup\$ @TimVermeulen: Well, I was too lazy to implement an insert method, therefore I just build a completely balanced tree, that might explain the difference. \$\endgroup\$
    – Martin R
    May 7, 2016 at 12:53
  • \$\begingroup\$ Interesting. I didn't really expect a balanced tree to perform better in traversal... \$\endgroup\$ May 7, 2016 at 13:10
  • \$\begingroup\$ I also implemented a red-black tree, and unfortunately, traversal was still relatively slow. \$\endgroup\$ May 7, 2016 at 13:16
  • \$\begingroup\$ great thinking, @MartinR. The "easy" way to traverse a tree is just recursively; I've been wondering how one could encapsulate a recursive traverse inside IteratorProtocol or generator. It's kind of confusing - you'd have to have "pauses", return that one, and "continue from that point" when called again for the next one. (A bit like programming a frame-based game engine, when you use coroutines to basically "wait" until the next frame.) \$\endgroup\$
    – Fattie
    Nov 1, 2016 at 20:41

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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