10. Trees

Written by Kelvin Lau

The tree is a data structure of profound importance. It is used in numerous facets of software development, such as:

• Representing hierarchical relationships.
• Managing sorted data.
• Facilitating fast lookup operations.

There are many types of trees, and they come in various shapes and sizes. In this chapter, you will learn the basics of using and implementing a tree.

Terminology

Many terms are associated with trees, and here are some you should know right off the bat.

Node

Like the linked list, trees are made up of nodes.

Each node can carry some data and keeps track of its children.

Parent and child

Trees are viewed starting from the top and branching towards the bottom, just like a real tree, only upside-down.

Every node (except for the topmost one) connects to exactly one node above it. That node is called a parent node. The nodes directly below and connected to it are called its child nodes. In a tree, every child has exactly one parent. That’s what makes a tree, well, a tree.

Root

The topmost node in the tree is called the root of the tree. It is the only node that has no parent:

Leaf

A node is a leaf if it has no children:

You will run into more terms later on, but this should be enough to get you started.

Implementation

Open up the starter playground for this chapter to get started. A tree consists of nodes, so your first task is to create a TreeNode class.

Create a new file named TreeNode.swift and write the following inside it:

public class TreeNode<T> {
  public var value: T
  public var children: [TreeNode] = []

  public init(_ value: T) {
    self.value = value
  }
}

Each node is responsible for a value and holds references to all its children using an array.

Note: Using a class type to represent TreeNode will mean losing value semantics. On the other hand, it makes creating references to nodes trivial, which you’ll use later on.

Next, add the following method inside the TreeNode class:

public func add(_ child: TreeNode) {
  children.append(child)
}

This method adds a child node to a node.

Time to give it a whirl. Head back to the playground page and write the following:

example(of: "creating a tree") {
  let beverages = TreeNode("Beverages")

  let hot = TreeNode("Hot")
  let cold = TreeNode("Cold")

  beverages.add(hot)
  beverages.add(cold)
}

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

Traversal algorithms

Iterating through linear collections such as arrays or linked lists is straightforward. Linear collections have a clear start and end:

Iterating through trees is a bit more complicated:

Should nodes on the left have precedence? How should the depth of a node relate to its precedence? Your traversal strategy depends on the problem that you’re trying to solve. There are multiple strategies for different trees and different problems. In the next section, you will look at depth-first traversal, a technique that starts at the root and visits nodes as deep as it can before backtracking.

Depth-first traversal

Write the following at the bottom of TreeNode.swift:

extension TreeNode {
  public func forEachDepthFirst(visit: (TreeNode) -> Void) {
    visit(self)
    children.forEach {
      $0.forEachDepthFirst(visit: visit)
    }
  }
}

This simple code uses recursion to process the next node.

You could use your own stack if you didn’t want your implementation to be recursive.

Time to test it out. Head back to the playground page and write the following:

func makeBeverageTree() -> TreeNode<String> {
  let tree = TreeNode("Beverages")

  let hot = TreeNode("hot")
  let cold = TreeNode("cold")

  let tea = TreeNode("tea")
  let coffee = TreeNode("coffee")
  let chocolate = TreeNode("cocoa")

  let blackTea = TreeNode("black")
  let greenTea = TreeNode("green")
  let chaiTea = TreeNode("chai")

  let soda = TreeNode("soda")
  let milk = TreeNode("milk")

  let gingerAle = TreeNode("ginger ale")
  let bitterLemon = TreeNode("bitter lemon")

  tree.add(hot)
  tree.add(cold)

  hot.add(tea)
  hot.add(coffee)
  hot.add(chocolate)

  cold.add(soda)
  cold.add(milk)

  tea.add(blackTea)
  tea.add(greenTea)
  tea.add(chaiTea)

  soda.add(gingerAle)
  soda.add(bitterLemon)

  return tree
}

This function creates the following tree:

Next, add this:

example(of: "depth-first traversal") {
  let tree = makeBeverageTree()
  tree.forEachDepthFirst { print($0.value) }
}

This code produces the following depth-first output:

---Example of: depth-first traversal---
Beverages
hot
tea
black
green
chai
coffee
cocoa
cold
soda
ginger ale
bitter lemon
milk

In the next section, you will look at level-order traversal, a technique that visits each node of the tree based on the depth of the nodes.

Level-order traversal

Write the following at the bottom of TreeNode.swift:

extension TreeNode {
  public func forEachLevelOrder(visit: (TreeNode) -> Void) {
    visit(self)
    var queue = Queue<TreeNode>()
    children.forEach { queue.enqueue($0) }
    while let node = queue.dequeue() {
      visit(node)
      node.children.forEach { queue.enqueue($0) }
    }
  }
}

forEachLevelOrder visits each of the nodes in level-order:

Note how you used a queue (not a stack) to ensure you visit nodes in the right level order. A simple recursion (which implicitly uses a stack) would not have worked!

Head back to the playground page and write the following:

example(of: "level-order traversal") {
  let tree = makeBeverageTree()
  tree.forEachLevelOrder { print($0.value) }
}

In the console, you will see the following output:

---Example of: level-order traversal---
Beverages
hot
cold
tea
coffee
cocoa
soda
milk
black
green
chai
ginger ale
bitter lemon

Search

You already have a method that iterates through all the nodes, so building a search algorithm shouldn’t take long. Write the following at the bottom of TreeNode.swift:

extension TreeNode where T: Equatable {
  public func search(_ value: T) -> TreeNode? {
    var result: TreeNode?
    forEachLevelOrder { node in
      if node.value == value {
        result = node
      }
    }
    return result
  }
}

Head back to the playground page to test your code. To save some time, copy the previous example and modify it to test the search method:

example(of: "searching for a node") {
  // tree from the last example

  if let searchResult1 = tree.search("ginger ale") {
    print("Found node: \(searchResult1.value)")
  }
  if let searchResult2 = tree.search("WKD Blue") {
    print(searchResult2.value)
  } else {
    print("Couldn't find WKD Blue")
  }
}

You will see the following console output:

---Example of: searching for a node---
Found node: ginger ale
Couldn't find WKD Blue

Here, you used your level-order traversal algorithm. Since this code visits all nodes, the last match will win if there are multiple matches. This instability means that you will get different objects back depending on what traversal you use.

Key points

• Trees share similarities to linked lists, but a tree node can link to many child nodes where linked-list nodes may only link to one successor node.
• Every tree node, except for the root node, has exactly one parent node.
• A root node has no parent nodes.
• Leaf nodes have no child nodes.
• Be comfortable with the tree terminology such as parent, child, leaf and root. Many of these terms are common tongue for fellow programmers and will help explain other tree structures.
• Traversals, such as depth-first and level-order traversals, aren’t specific to the general tree. They work on other kinds of trees, although their implementation will be slightly different based on how the tree is structured.