Tim discusses using Swift enumeration indirect keywords to build binary tree nodes. Read on to learn more about how reference types and value semantics combine…

Today we’re writing a simple binary search tree in Swift. While binary search trees seem to be very powerful in theory, their performance is rather disappointing in practice. Other, more advanced tree-shaped data structures such as red-black trees and B-trees solve some of the practical problems that binary search trees have, and are subsequently much more useful in performance-critical code. Nevertheless, learning about binary search trees is the first step to getting more familiar with this class of data structures.

Probably the simplest way to represent a binary tree in Swift is by using an indirect enum, like so:

enum BinarySearchTree<Element: Comparable> {
    case empty
    indirect case node(left: BinarySearchTree, value: Element, right: BinarySearchTree)
}

While indirect enums use reference types under the hood, they have value semantics by default. So that’s nice.

From all the standard library’s protocols, SetAlgebra and BidirectionalCollection both seem a good fit for our BinarySearchTree type: SetAlgebra for inserting and removing elements, and BidirectionalCollection for traversing the tree (both forwards and backwards, hence the name). However, for the purpose of this post, we’ll stick to only a couple basic methods.

Let’s start with insertion. Because of the way enums with associated values work, it’s easiest to implement the insert method using an under-the-hood inserting method:

extension BinarySearchTree {
    mutating func insert(_ element: Element) {
        self = inserting(element)
    }
    
    private func inserting(_ element: Element) -> BinarySearchTree {
        switch self {
        case .empty:
            // the tree is empty, so inserting an element results in a tree containing only that element
            return .node(.empty, element, .empty)
            
        case .node(_, element, _):
            // the element is already present in the tree
            return self
            
        case let .node(left, value, right) where element < value:
            // the element should be inserted into the left subtree
            return .node(left.inserting(element), value, right)
            
        case let .node(left, value, right):
            // the element should be inserted into the right subtree
            return .node(left, value, right.inserting(element))
        }
    }
}

Now we can insert values into a tree, but we can’t read them in any way. So let’s add a contains method as well:

extension BinarySearchTree {
    func contains(_ element: Element) -> Bool {
        switch self {
        case .empty:
            // an empty tree obviously doesn't contain any elements!
            return false
            
        case .node(_, element, _):
            // the element is equal to this node's value
            return true
            
        case let .node(left, value, _) where element < value:
            // if the element is present in the tree, it must be in the left subtree
            return left.contains(element)
            
        case let .node(_, _, right):
            // if the element is present in the tree, it must be in the right subtree
            return right.contains(element)
        }
    }
}

Let’s try it out!

var tree = BinarySearchTree.empty

tree.contains(5) // => false
tree.insert(5)
tree.contains(5) // => true
tree.insert(3)
tree.contains(3) // => true
tree.contains(5) // => true

Looks good. We can’t actually remove elements, though, but that is a real pain to implement and it’s out of scope for this post.

Finally, to iterate over a binary search tree, we need to implement the IteratorProtocol protocol. We’ll use an in-order traversal algorithm:

extension BinarySearchTree: Sequence {
    func makeIterator() -> BinarySearchTreeIterator<Element> {
        return BinarySearchTreeIterator(self)
    }
}

struct BinarySearchTreeIterator<Element: Comparable>: IteratorProtocol {
    var node: BinarySearchTree
    var stack: [(Element, BinarySearchTree)] = []
    
    init(_ node: BinarySearchTree) {
        self.node = node
    }
    
    public mutating func next() -> Element? {
        while case let .node(left, value, right) = node {
            stack.append((value, right))
            node = left
        }
        
        guard let (element, node) = stack.popLast() else { return nil }
        
        self.node = node
        return element
    }
}

Now we can do all kinds of sequence-y stuff with trees, like so:

var tree = BinarySearchTree.empty
[5, 2, 4, 8, 3, 2].forEach { tree.insert($0) }

for element in tree {
    print(element, terminator: " ") // => 2 3 4 5 8
}

tree.reduce(0, +)                                  // => 22
tree.lazy.map(String.init).joined(separator: ", ") // => 2, 3, 4, 5, 8

And the list goes on.

That’s it for now! If you enjoy this kind of stuff, make sure to check out Károly Lőrentey’s brand new book Optimizing Collections in which he in much detail goes through implementing several data structures in Swift, focusing on performance. As of writing this post, it’s 25% off.