Swift Binary Search Tree With Examples [Latest]

Binary Search Tree

A Binary Search Tree in Swift is a Binary Tree which must follow the following rules:

All nodes in the left subtree must have lesser value than the value in the root node.
All node in the right subtree must have greater value than the value in the root node.
Left and Right subtrees of root node should follow the above rules recursively.

They typically provide OLogN time for insertion and lookup.

Swift Binary Search Tree Implementation

Let’s define the Node for a tree.

class Node<T> {
    var data: T
    var leftNode: Node?
    var rightNode: Node?
    init(data: T,
         leftNode: Node? = nil,
         rightNode: Node? = nil) {
        self.data = data
        self.leftNode = leftNode
        self.rightNode = rightNode
    }
}

Let’s define the BinarySearchTree class in Swift. Here we’ll define functions to insert Nodes into the tree taking care of the BST rules:

class BST<T: Comparable & CustomStringConvertible> {
    private var rootNode: Node<T>?
    func addNode(_ value: T) {
        let node = Node(data: value)
        if let rootNode = self.rootNode {
            self.insert(rootNode, node)
        } else {
            self.rootNode = node
        }
    }
    private func insert(_ root: Node<T>, _ node: Node<T>) {
        if root.data > node.data {
            if let leftNode = root.leftNode {
                self.insert(leftNode, node)
            } else {
                root.leftNode = node
            }
        } else {
            if let rightNode = root.rightNode {
                self.insert(rightNode, node)
            } else {
                root.rightNode = node
            }
        }
    }
    func printTree() {
        self.inorder(self.rootNode)
    }
    private func inorder(_ node: Node<T>?) {
        guard let _ = node else { return }
        self.inorder(node?.leftNode)
        print("(node!.data)", terminator: " ")
        self.inorder(node?.rightNode)
    }
}

We’ve assigned the Comparable and CustomStringConvertible protocols in order to compare the values of the Nodes and print the formatted data respectively.

The inorder function prints the left subtree followed by the current node value, followed by right subtree

Now let’s add elements into the tree.

let numberList : Array<Int> = [8, 2, 10, 9, 11, 1, 7]
var root = BST<Int>()
for number in numberList {
    root.addNode(number)
}
//should print sorted tree
root.printTree()

The output of the tree is:

As you can see the values in the BST are set in a sorted order.

Searching a value in the tree

extension BST {
    func search(value: T) {
        self.search(self.rootNode, value)
    }
    private func search(_ root: Node<T>?, _ element: T) {
        guard let rootNode = root else {
            print("NODE is Not available : (element)")
            return
        }
        print("current root node (rootNode.data)")
        if element > rootNode.data {
            self.search(rootNode.rightNode, element)
        } else if element < rootNode.data {
            self.search(rootNode.leftNode, element)
        } else {
            print("NODE VALUE AVAILABLE : (rootNode.data)")
        }
    }
}

We’ve created a search function in the extension.
In this, we simply check the node value and based on the comparison, search it recursively in the left or right subtree.

The output of two sample runs is:

Deleting a Node

Implementation of Deleting a Node in a BST is a little more tricky.

After the node is deleted, we need to rearrange the BST so that it stays sorted.
Use the following rule for deletion:

After removing a node, we replace the node with either its biggest child on the left or its smallest child on the right subtree.

This means we need to first create functions in our Node class for the minimum and maximum node of the tree.
The following illustration contains the updated class Node.

And now we create another extension of the BST class for the deletion function which works recursively:

extension BST{
    func deleteKey(value: T)
    {
    rootNode = deleteRec(rootNode, value)
    }
    /* A recursive function to insert a new key in BST */
    func deleteRec(_ root: Node<T>?, _ key: T) -> Node<T>?
    {
    /* Base Case: If the tree is empty */
        if  root == nil{
        return root
        }
        if key < (root?.data)! {
        root?.leftNode = deleteRec(root?.leftNode, key)
        }
        else if key > (root?.data)!{
            root?.rightNode = deleteRec(root?.rightNode, key)
            }
    else
    {
        if root?.leftNode == nil{
            return root?.rightNode
        }
        else if root?.rightNode == nil{
            return root?.leftNode
        }
    root?.data = (minValue(root?.rightNode))!
    root?.rightNode = deleteRec(root?.rightNode, (root?.data)!)
    }
    return root;
    }
    public func minValue(_ node: Node<T>?) -> T? {
        var tempNode = node
        while let next = tempNode?.leftNode {
            tempNode = next
        }
        return tempNode?.data
    }
}

The output when deleting the first key is:

That’s all for Swift Binary Search Tree implementation.