Swift Algorithm Club: Swift Tree Data Structure

The Swift Algorithm Club is an open source project to implement popular algorithms and data structures in Swift.

Every month, Chris Pilcher and I will write a tutorial showing you how to make a cool algorithm or data structure from the project.

This series will be a great way to learn more about algorithms and data structures, and to level up your Swift skills along the way.

In this first tutorial, you’ll learn how to implement a Swift Tree data structure. This is one of the most common and useful data structures, and is a great way to get started!

Tree Data Structure

The easiest way to understand the tree data structure is is through a picture:

The above diagram shows a tree with 5 levels. The root is level 0, and as you down the depth of the tree, the level increases by 1.

Trees can help you solve many important problems, including:

• representing a hierarchical relationship between objects
• making searches quick and efficient
• providing a sorted list of data
• powering prefix matching in text fields

Terminology

First, let’s cover some important terminology you should understand about trees.

Root

The root of a tree refers to the 0th level’s node of a tree. You can also think of it as the entry point to your tree data structure.

Node

A node is a block of data in the tree structure. The data that a node contains depends on the type of tree you’re making. The root is also a node.

Leaf

Sometimes referred as a terminal node, a leaf is simply a node with no children. For example, if the Node object had leftChild and rightChild as nil, then you would refer that node as a leaf.

Tree Implementation in Swift

In this section, you’ll implement a general-purpose tree. This is a fancy way of saying a tree without any kind of restrictions (like how many children each node may have, or the order of nodes).

Remember that a tree is made up of nodes. So to start, let’s create a basic node class. Create a new Swift playground and add the following empty class:

class Node {
 
}

Value

Of course, a node isn’t much use without a value associated with it.

For simplicity, you’ll specialize this tree to manage string data. Update your current implementation of the Node to the following:

class Node {
  var value: String
 
  init(value: String) {
    self.value = value
  }
}

You’ve declared a property named value of type String. You also declare an initializer, which is required for initializing all non-optional stored properties for your class.

Children

In addition to a value, each node needs to have a list of children.

Update your class definition to the following:

class Node {
  var value: String
  var children: [Node] = [] // add the children property
 
  init(value: String) {
    self.value = value
  }
}

You simply declare children as an array of nodes. Each child represents a node that is 1 level deeper than the current node.

Parent

Sometimes it’s handy for each node to have a link to its parent node as well. Children are the nodes below a given node; the parent is the node above. A node may only have one parent, but can have multiple children.

Update the implementation of your Node class to the following:

class Node {
  var value: String
  var children: [Node] = []
  weak var parent: Node? // add the parent property
 
  init(value: String) {
    self.value = value
  }
}

Note that you’ve made parent an optional. This is because not all nodes have parents – such as the root node in a tree.

Insertion

To handle insertion to your tree, you’ll declare a addChild() method in your Node class. Update the implementation of your class to the following:

class Node {
  var value: String
  var children: [Node] = []
  weak var parent: Node?
 
  init(value: String) {
    self.value = value
  }
 
  func addChild(node: Node) {
    children.append(node)
    node.parent = self
  }
}

It’s best to understand how addChild() works by using it on a live playground. Outside the implementation of your class, write the following into your playground:

let beverages = Node(value: "beverages")
 
let hotBeverages = Node(value: "hot")
let coldBeverages = Node(value: "cold")
 
beverages.addChild(hotBeverages)
beverages.addChild(coldBeverages)

Hierarchical structures are natural candidates for tree structures, so here you’ve defined 3 different nodes and organized them into a logical hierarchy. This corresponds to the following structure:

Challenge: Beverage City

Ready for a quick test of knowledge?

Try writing the code to extend your tree to match the following diagram:

The solution is provided in the spoiler section down below, but try it yourself first!

let beverages = Node(value: "beverages")
 
let hotBeverage = Node(value: "hot")
let coldBeverage = Node(value: "cold")
 
let tea = Node(value: "tea")
let coffee = Node(value: "coffee")
let cocoa = Node(value: "cocoa")
 
let blackTea = Node(value: "black")
let greenTea = Node(value: "green")
let chaiTea = Node(value: "chai")
 
let soda = Node(value: "soda")
let milk = Node(value: "milk")
 
let gingerAle = Node(value: "ginger ale")
let bitterLemon = Node(value: "bitter lemon")
 
beverages.addChild(hotBeverage)
beverages.addChild(coldBeverage)
 
hotBeverage.addChild(tea)
hotBeverage.addChild(coffee)
hotBeverage.addChild(cocoa)
 
coldBeverage.addChild(soda)
coldBeverage.addChild(milk)
 
tea.addChild(blackTea)
tea.addChild(greenTea)
tea.addChild(chaiTea)
 
soda.addChild(gingerAle)
soda.addChild(bitterLemon)

Printing Trees

Verifying a large tree structure can be hard without any console logging. After defining your tree structure, try to log your result in the console by printing the tree:

print(beverages) // <- try to print it!

You can bring up the console by pressing the following keys in combination: Command-Shift-Y. You should see the name of your class printed onto the console.

Node

How silly! Unfortunately the compiler doesn’t know the best print your custom Swift object, unless you tell it.

To aid the compiler, you’ll need to make Node adopt the CustomStringConvertible protocol. To do this, add the following just below the implementation of your Node class:

// 1
extension Node: CustomStringConvertible {
  // 2
  var description: String {
    // 3
    var text = "\(value)"
 
   // 4
    if !children.isEmpty {
      text += " {" + children.map { $0.description }.joinWithSeparator(", ") + "} "
    }
    return text
  }
}

This code is relatively straight forward:

1. You’ve declared an extension to your Node class, and you’ve adopted the CustomStringConvertible protocol. This protocol expects you to implement a computed property with the name description, with the String type.
2. You’ve declared the description property. This is a computed property, a read only property that returns a String.
3. You’ve declared a text variable. This will hold the entire string. For now, you’ve given it the current node’s value.
4. In addition to printing the current value of the node, you’ll also need to print the children, children of children, and so on. To do so, you’ll recursively append your children’s description, whilst adding some braces to give the string some context in regards to the structure of the children.

Now, when you call the print your Node classes, you’ll get a nice representation of your tree structure like this:

"beverages {hot {tea {black, green, chai} , coffee, cocoa} , cold {soda {ginger ale, bitter lemon} , milk} } \n"

if !children.isEmpty {
  text += " {"
  for child in children {
    text += child.description + ", "
  }
  text += "} "
}

Search

The general-purpose tree shown here is great for describing hierarchical data, but it really depends on your application in regards to what kind of extra functionality it needs to have. For example, you could use the Node class to determine if the tree contains a particular value.

To facilitate a search algorithm for this general-purpose tree, add the following extension at the bottom of your playground file:

extension Node {
  // 1
  func search(value: String) -> Node? {
    // 2
    if value == self.value {
      return self
    }
    // 3
    for child in children {
      if let found = child.search(value) {
        return found
      }
    }
    // 4
    return nil
  }
}

The code here is relatively straightforward:

1. The goal of this method is to search if a value exists in the tree. If it does, return the node associated with the value. If it does not exist, you’ll return a nil.
2. This is the case where you’ve found the value. You’ll return the self, which is the current node.
3. In this loop, you cycle through the children array. You’ll call each child’s search method, which will recursively iterate through all the children. If any of the nodes have a match, your if let statement will evaluate true and return the node.
4. You’ll return nil here to signify that you couldn’t find a match.

Let’s give our search method a try! At the bottom of your playground file, write the following:

beverages.search("cocoa") // returns the "cocoa" node
beverages.search("chai") // returns the "chai" node
beverages.search("bubbly") // returns nil

What About Different Types?

Nice work so far! You’ve learned how to implement a general-purpose tree that stores String values. You’ve defined a nice way to print your tree into the console, and also provided searching capabilities to your Node class.

Trees are a great way to lay out your hierarchical structure of strings, but what if you wanted to store integers instead?

You could modify the Node class to take in an Int:

class Node {
  var value: Int
 
  // ...
}

but then your old implementation that accepts String value is lost. Ideally, you’d want to create a Node class that could accept all types of objects, whether it is an Int, Double, Float, or even a custom class of your own. To facilitate the generic usage of your Node class, you’ll have to dive in the world of generics!

Generics

The idea of generics is to abstract away the type requirements from algorithm and data structures. This allows you to keep the idea generalized and reusable. Whether an object would behave well in a tree (or any other data structure) should not be whether it is an Int or a String, but rather something more intrinsic; In the context of trees, any type that behaves well in a hierarchy are good candidates to be used in a tree.

Time to make some breaking changes! Update the implementation of your Node class to the following:

// 1. 
class Node<T> {
  // 2. 
  var value: T
  weak var parent: Node?
  // 3.
  var children: [Node] = []
 
  // 4.
  init(value: T) {
    self.value = value
  }
 
  // 5. 
  func addChild(node: Node) {
    children.append(node)
    node.parent = self
  }
}

Right away, you should see some compiler errors. But fear not, you’ll clear those errors as you finish the generalizing your Node implementation for all types. Here’s what you’ve done here:

1. You’ve changed the declaration of the Node class to take a generic type T. The <> syntax around T is what alerts the compiler that your intention of using generics.
2. Your goal is to allow the Node class to take in values of any type, so you’ll constrain your value property to be type T rather than Int or String.
3. For the same reason as the other points, you’ll now declare that your class has children of type T.
4. You’ve also updated your initializer to take any type.
5. You’ve updated your addChild method to take in Node objects of any type matching the current type of Node

So far so good. Next, find the extension that contains the search method and update it to use generics:

// 1. 
extension Node where T: Equatable {
  // 2. 
  func search(value: T) -> Node? {
    if value == self.value {
      return self
    }
    for child in children {
      if let found = child.search(value) {
        return found
      }
    }
    return nil
  }
}

You’ve made two changes here:

1. You’ve introduced a constraint for this extension so that any type must be Equatable before it can utilize the search method.
2. You’ve updated the value parameter to be of generic type.

Your code should compile now, so let’s test this out! At the bottom of your playground file, add the following code to verify that your generic tree is working:

let number = Node(value: 5)

Congratulations, you’ve just create a general-purpose tree that works for all types of objects!

Other Trees

You’ve created a very basic tree here, but there are many different ways to construct trees. For example:

• Sometimes you don’t need to have a parent property at all.
• Maybe you only need to give each node a maximum of two children – such a tree is called a binary tree.
• A very common type of tree is the binary search tree (or BST), a stricter version of a binary tree where the nodes are ordered in a particular way to speed up searches.

To learn more about these kinds of trees and more, check out this list of articles in the Swift Algorithm Club repo:

• AVL Tree
• B-Tree
• Binary Search Tree
• Binary Tree
• Minimum Spanning Tree (unweighted)
• Radix Tree
• Red-Black Tree
• Segment Tree
• Threaded Binary Tree
• Tries
• Union-Find

Where To Go From Here?

I hope you enjoyed this tutorial on making a Swift Tree data structure!

It’s in your best interest to know about algorithms and data structures – they’re solutions to many real world problems, and are frequently asked as interview questions. Plus it’s fun!

So stay tuned for many more tutorials from the Swift Algorithm club in the future. In the meantime, if you have any questions on implementing trees in Swift, please join the forum discussion below!

The post Swift Algorithm Club: Swift Tree Data Structure appeared first on Ray Wenderlich.