Swift Algorithm Club: Swift Tree Data Structure

Learn how to implement a Swift Tree data structure through this hands-on tutorial! By Kelvin Lau.

Leave a rating/review
Save for later
Share
You are currently viewing page 2 of 2 of this article. Click here to view the first page.

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: 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 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(value: "cocoa") // returns the "cocoa" node
beverages.search(value: "chai") // returns the "chai" node
beverages.search(value: "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 a 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 algorithms 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 is a good candidate 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 add(child: Node) {
    children.append(child)
    child.parent = self
  }
}

Here's what you've done:

  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 to 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 add(child:) 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: 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 a 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 created 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:

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!

Note: The Swift Algorithm Club is always looking for more contributors. If you've got an interesting data structure, algorithm, or even an interview question to share, don't hesitate to contribute! To learn more about the contribution process, check out our Join the Swift Algorithm Club article.

If you enjoyed what you learned in this tutorial, why not check out our Data Structures and Algorithms in Swift book, available on our store?

In Data Structures and Algorithms in Swift, you’ll learn how to implement the most popular and useful data structures and when and why you should use one particular datastructure or algorithm over another. This set of basic data structures and algorithms will serve as an excellent foundation for building more complex and special-purpose constructs.

As well, the high-level expressiveness of Swift makes it an ideal choice for learning these core concepts without sacrificing performance.

  • You’ll start with the fundamental structures of linked lists, queues and stacks, and see how to implement them in a highly Swift-like way.
  • Move on to working with various types of trees, including general purpose trees, binary trees, AVL trees, binary search trees and tries.
  • Go beyond bubble and insertion sort with better-performing algorithms, including mergesort, radix sort, heap sort and quicksort.
  • Learn how to construct directed, non-directed and weighted graphs to represent many real-world models, and traverse graphs and trees efficiently with breadth-first, depth-first, Dijkstra’s and Prim’s algorithms to solve problems such as finding the shortest path or lowest cost in a network.
  • And much, much more!

By the end of this book, you’ll have hands-on experience solving common issues with data structures and algorithms — and you’ll be well on your way to developing your own efficient and useful implementations.