Implement Tree with Swift

Let’s start it with the definition of tree. What is Tree in programming? Tree is a data structure that contains a hierarchical relation between its elements. Tree has so many types, but in this topic, we only talk about Tree in general. For other types, we will talk about it in a specific topic later.

We can see that the node is an element that has branch/es that connected with other elements. The leaf is an element that doesn’t have any branch. The root is the highest level hierarchy node.

And now let’s code it.

class Node<ValueType: Equatable> {
    weak var parent: TreeNode?
    var data: ValueType
    var children: [TreeNode] = []
    
    init(data: ValueType) {
        self.data = data
    }
}

Each node will contain parent, data, and childer (array of nodes that become the children). For the parent variable, we set it as a weak variable to avoid the `memory cycle problem`. If you don't understand it, you can read about Automatic Reference Counting (ARC).

For now, we only make 2 methods in this function, those are `add(child:)` and `search(element:)`.

class TreeNode<ValueType: Equatable> {
    weak var parent: TreeNode?
    var data: ValueType
    var children: [TreeNode] = []
    
    init(data: ValueType) {
        self.data = data
    }
    
    func add(child: TreeNode) {
        child.parent = self
        children.append(child)
    }
    
    func search(element: ValueType) -> TreeNode? {
        if element == data {
            return self
        }
        
        for child in children {
            if let result = child.search(element: element) {
                return result
            }
        }
        
        return nil
    }
}

And we need a method to print the hierarchy in the TreeNode to see does our code run properly or not. Here’s the implementation for the print the TreeNode Hierarchy:

func printTreeHierarchy() {
    print(getNodeHierarchy().joined(separator: "\n"))
}
private func getNodeHierarchy(level: Int = 0) -> [String] {
    var spacing: String = ""
    
    for _ in 0..<level {
        spacing += "\t"
    }
    
    var results: [String] = [spacing + "\(data)"]
    
    for child in children {
        results += child.getNodeHierarchy(level: level + 1)
    }
    
    return results
}

Now, we’ll make implementation for a tree like this:

Here’s the code:

let techTeam: TreeNode<String> = TreeNode<String>(data: "Technology Team")
let frontEndTeam: TreeNode<String> = TreeNode<String>(data: "Front End Team")

techTeam.add(child: frontEndTeam)
techTeam.add(child: TreeNode<String>(data: "Back End Team"))
techTeam.add(child: TreeNode<String>(data: "Data Team"))

frontEndTeam.add(child: TreeNode<String>(data: "iOS Team"))
frontEndTeam.add(child: TreeNode<String>(data: "Android Team"))
frontEndTeam.add(child: TreeNode<String>(data: "Web Team"))

techTeam.printTreeHierarchy()

And the result will be:

Technology Team
    Front End Team
        iOS Team
        Android Team
        Web Team
    Back End Team
    Data Team

So, here’s the explanation about the basics of Tree. We will talk about other types of trees in another specific topic.