11. Tree Challenges

Written by Kelvin Lau

Challenge 1: Print a tree in level order

Print all the values in a tree in an order based on their level. Nodes in the same level should be printed on the same line. For example, consider the following tree:

Your algorithm should print the following:

15
1 17 20
1 5 0 2 5 7

Hint: Consider using a Queue included for you in the Sources folder of the starter playground.

Challenge 2: Parents and ownership

Consider the original definition of a tree node:

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

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

How can you modify this definition to include a parent? What considerations should you make about ownership?

Solutions

Solution to Challenge 1

A straightforward way to print the nodes in level-order is to leverage the level-order traversal using a Queue data structure. The tricky bit is determining when a newline should occur. Here’s the solution:

func printEachLevel<T>(for tree: TreeNode<T>) {
  // 1
  var queue = Queue<TreeNode<T>>()
  var nodesLeftInCurrentLevel = 0
  queue.enqueue(tree)
  
  // 2
  while !queue.isEmpty {
  
    // 3
    nodesLeftInCurrentLevel = queue.count
    
    // 4
    while nodesLeftInCurrentLevel > 0 {
      guard let node = queue.dequeue() else { break }
      print("\(node.value) ", terminator: "")
      node.children.forEach { queue.enqueue($0) }
      nodesLeftInCurrentLevel -= 1
    }
    
    // 5
    print()
  }
}

1. You begin by initializing a Queue data structure to facilitate the level-order traversal. You also create nodesLeftInCurrentLevel to keep track of the number of nodes you’ll need to work on before you print a new line.

2. Your level-order traversal continues until your queue is empty.

3. Inside the first while loop, you begin by setting nodesLeftInCurrentLevel to the current elements in the queue.

4. Using another while loop, you dequeue the first nodesLeftInCurrentLevel number of elements from the queue. Every element you dequeue is printed out without establishing a new line. You also enqueue all the children of the node.

5. At this point, you generate the new line using print(). In the next iteration, nodesLeftInCurrentLevel will be updated with the count of the queue, representing the number of children from the previous iteration.

This algorithm has a time complexity of O(n). Since you initialize the Queue data structure as an intermediary container, this algorithm also uses O(n) space.

Solution to Challenge 2

You can add a property parent to the TreeNode like so:

public class TreeNode<T> {

  public weak var parent: TreeNode?

  // etc...
}

An optional type is used since the root node does not have a parent. It is given weak ownership to conveniently avoid reference cycles. By convention, a node has a strong ownership relationship with its children but a weak non-ownership relationship with its parent. Continuing the linked list analogy, having nodes with a parent is analogous to a doubly linked list. There is more book keeping overhead to worry about, but it allows quick upward traversal of the tree.