Published on: May 10, 2020

Tree in Swift

Tree data structure is something I am sure every programmer has implemented at least once (though more for interviews than in job 😁). Nevertheless let's see how to represent a binary tree in Swift.

A binary tree is a specific form of tree data structure where each node has at most two children which are referred to as the left child and the right child. [Ref]

You probably figured out from the definition that left and right children are itself Trees, so let's start with a class which is a Ref type (with a value type like structure we cannot represent a Tree, which means an infinite size memory)

class TreeClass<T> {
    var data: T
    var left: TreeClass<T>?
    var right: TreeClass<T>?

    init(data: T, left: TreeClass<T>?, right: TreeClass<T>?) {
        self.data = data
        self.left = left
        self.right = right
    }
}

Now represent a tree with root (3) left child (1) and right child (2)

let tree11 = TreeC(data: 1, left: nil, right: nil)
let tree22 = TreeC(data: 2, left: nil, right: nil)
let tree33 = TreeC(data: 3, left: tree11, right: tree22)

Let's write a preorder traversal algorithm

func preorder<T>(tree: TreeC<T>?) {
    guard let t = tree else {
        return
    }
    print(t.data)
    preorder(tree: t.left)
    preorder(tree: t.right)
}

In Swift there is another (better) way to represent recursive data structures like Tree and that is called Indirect enum

With it an enum case can refer to itself which essentially is possible as compiler internally make it like a reference type. (for normal enums which are value type, internally compiler treat them as C unions )

indirect enum Tree<T> {
  case Nil
  case Node(value: T ,left: Tree<T> ,right: Tree<T>)
}

Much better, right?

Same tree representation

let tree1 = Tree.Node(value: 1, left: Tree.Nil, right: Tree.Nil)
let tree2 = Tree.Node(value: 2, left: Tree.Nil, right: Tree.Nil)
let tree3 = Tree.Node(value: 3, left: tree1, right: tree2)

And preorder algorithm

func preorder<T>(tree:Tree<T>) {
    switch tree {
    case .Nil:
        break;
    case .Node(let value, let left, let right):
        print(value)
        preorder(tree: left)
        preorder(tree: right)
    }
}

Swift is evolving very fast and has so many features that sometime we miss to use the best feature available for a given problem.

Tagged with: