Swift Algorithm Club: Swift Trie Data Structure

Learn how to implement the trie data structure in Swift – a data structure that is extremely handy for prefix-matching in the English language. By Kelvin Lau.

Leave a rating/review
Save for later
Share

Swift Algorithm Club: Swift Trie Data Structure

15 mins

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

Every month, Chris Pilcher and I feature a cool data structure or algorithm from the club in a tutorial on this site. If your want to learn more about algorithms and data structures, follow along with us!

In this tutorial, you’ll learn how to implement a Swift Trie data structure. No, this is not “Tree” spelled wrong; this is actually a different data structure!

This algorithm was first implemented by Christian Encarnacion, and is now refactored for tutorial format.

Note: New to the Swift Algorithm Club? Check out our getting started post first.

Getting Started

Tries are n-ary trees in which characters are stored at each node. In addition to being a popular topic amongst interview questions, Tries are also a key data structure that facilitates efficient prefix matching for the English language:

Example of “Cat”, “Cut”, “Cute”, “To”, “B” strings stored in a Trie.

Example of “Cat”, “Cut”, “Cute”, “To”, “B” strings stored in a Trie.

Why a Trie?

Tries are very useful for certain situations. In addition to be great for storing the English language, a Trie can also be a substitute for a hash table, with the following advantages:

• Looking up values typically have a better worst-case time complexity.
• Unlike a hash table, a Trie does not need to worry about key collisions.
• Doesn’t require a hashing algorithm to guarantee a unique path to elements.
• Trie structures can be alphabetically ordered.

In this tutorial, you’ll focus on Trie’s application for storing the English language.

Trie Implementation

Just like other trees, a Trie is made up of nodes. Your implementation will consist of a `TrieNode` class and a `Trie` class. Each `TrieNode` will represent a character of a word. For instance, the word “cute” will be represented by the following series of nodes: `c -> u -> t -> e`. The `Trie` class will manage the insertion logic and keep a reference to the nodes.

Open up a Swift playground to begin!

TrieNode

You’ll start off by implementing a simple `TrieNode` class. Write the following into the playground:

```class TrieNode<T: Hashable> {
var value: T?
weak var parent: TrieNode?
var children: [T: TrieNode] = [:]

init(value: T? = nil, parent: TrieNode? = nil) {
self.value = value
self.parent = parent
}
}
```

This is a generic `TrieNode` class. It stores a value (i.e. the character) and has a reference to its parent and children. There are two things to point out:

• The `parent` property is `weak` to prevent reference cycles. Having a reference to the parent is necessary for `remove` operations on the Trie.
• The value stored in `TrieNode` must conform to the `Hashable` protocol. This is because you will be using the value as a key in the `children` dictionary – and anything that is a key in a Swift dictionary must conform to `Hashable`. You will be using `Character` for the value, which conforms to `Hashable`, so you are set.

To facilitate the adding of new nodes, add the following method inside the `Node` class:

```func add(child: T) {
// 1
guard children[child] == nil else { return }

// 2
children[child] = TrieNode(value: child, parent: self)
}
```

Adding a child is a 2-stage process:

1. Make sure the child does not already exist in the dictionary of children. If it does, return.
2. Create a new node for the new value, and add it to the children dictionary of the current node.

With that, you’ve got a fairly familiar node object common to many trees. It’s still missing a component to be useful for a Trie, but you’ll handle that later :]

Trie

Your Trie class will be managing the nodes. Write the following at the bottom of your playground file:

```class Trie {
fileprivate let root: TrieNode<Character>

init() {
root = TrieNode<Character>()
}
}
```

This sets the foundation for your Trie. You declare a `root` property that keeps a reference to the root node of your Trie. Since you’re implementing a Trie for the English language, you’ll use nodes of type `Character`. The `init` method simply initializes the `root` property with an empty `TrieNode`.

Typealiasing

Before continuing on with implementing the rest of the Trie, update the `Trie` class to the following:

```class Trie {
typealias Node = TrieNode<Character>
fileprivate let root: Node

init() {
root = Node()
}
}
```

You’ve added a `Node` typealias. While this is functionally identical to your previous version, this allows you to refer to the `TrieNode` types as `Node`. In additional shortening the syntax, you also make the code more robust; If you ever wanted the node to represent something else other than a `Character`, changing just the `typealias` would propagate the type to everything else!

With that done, it’s time to implement the methods that make up the Trie.

Insertion

The `Trie` class manages the operations on the Trie. When implementing the insertion method, remember that a Trie is efficient because it always tries (pun intended) to reuse existing nodes to complete a sequence.

As an example, the two words “Cut” and “Cute” should be represented using 4 nodes, since both words share the same “Cut” prefix.

Add the following code below the `Trie` class:

```extension Trie {
func insert(word: String) {
// 1
guard !word.isEmpty else { return }

// 2
var currentNode = root

// 3
let characters = Array(word.lowercased().characters)
var currentIndex = 0

// ... more to come!
}
}
```

You’ve implemented the `insert` method in an extension. Here’s what you’ve written so far:

1. Check if the string is empty. If it is, there’s nothing to insert!
2. Create a reference to the root node. You’ll use this to iterate through the Trie nodes.
3. A word in the Trie is represented by a chain of nodes, where each node represents a character of the word (Ex: `c -> u -> t -> e` for “cute”). Since you’ll be inserting character by character, turning the word into an array will easily allow you to keep track of the characters during insertion.

Now that you’ve got the pieces ready, you’re ready to perform some pointer arithmetic! Add the following to the end of the `insert` method:

```while currentIndex < characters.count {
// 1
let character = characters[currentIndex]

// 2
if let child = currentNode.children[character] {
currentNode = child
} else {
// 3
currentNode = currentNode.children[character]!
}

// 4
currentIndex += 1

// more to come!
}
```

This code is relatively straight forward:

1. Get ahold of the character you need to insert into the Trie.
2. Check if the character you're trying to insert exists within the current node's `children` dictionary. If it exists, you'll simply move the `currentNode` reference to the next node. There's no need to insert the character because it's already there!
3. If execution proceeds to the else block, it means the character needs to be inserted. You'll add the character into the current `children` dictionary. Afterwards, you'll move the `currentNode` reference to the new node.
4. Add 1 to the `currentIndex` property to keep track of the next character you need to insert.

Contributors

Author

Final Pass Editor