Tries in Kotlin: Building an Autocomplete System

A Trie, also known as a Prefix Tree, is a specialized tree-based data structure used to store a set of strings. Unlike a Hash Map, where keys are stored as a whole, a Trie stores keys character by character.

This structure is incredibly efficient for features we use every day, such as search engine autocomplete and spell checkers.

---

1. The Trie Node

In a Trie, each node represents a character. It contains a map of its children (the characters that can follow it) and a boolean flag to indicate if it marks the completion of a word.

data class TrieNode(
    val children: MutableMap<Char, TrieNode> = sortedMapOf<Char, TrieNode>(),
    var isEndOfWord: Boolean = false
)

2. Core Operations: Insert and Search

Insertion

To insert a word, we iterate through its characters. If a character doesn’t exist as a child of the current node, we create a new node. Finally, we mark the last node as isEndOfWord.

fun insert(word: String, root: TrieNode) {
    var current = root
    for (char in word) {
        current = current.children.getOrPut(char) { TrieNode() }
    }
    current.isEndOfWord = true
}

Search

Searching follows the same path. If we successfully navigate through all characters and the final node reached is marked as an end of a word, the word exists in our Trie.

fun search(word: String, root: TrieNode): Boolean {
    var currNode = root

    for(i in word) {
        if(currNode.children[i] != null) {
            currNode = currNode.children[i]!!
        } else {
            return false
        }
    }
    return currNode.isEndOfWord
}

3. Autocomplete and Prefix Suggestions

fun getSuggestions(prefix: String, root: TrieNode): List<String> {
    var current = root

    // 1. Move to the node representing the end of the prefix
    for (char in prefix) {
        val node = current.children[char] ?: return emptyList()
        current = node
    }

    // 2. Collect all words from this node downward
    val results = mutableListOf<String>()
    collectWords(current, StringBuilder(prefix), results)
    return results
}

private fun collectWords(node: TrieNode, currentWord: StringBuilder, results: MutableList<String>) {
    if (node.isEndOfWord) {
        results.add(currentWord.toString())
    }

    for ((char, childNode) in node.children) {
        currentWord.append(char)
        collectWords(childNode, currentWord, results)
        currentWord.deleteCharAt(currentWord.length - 1) // Backtrack
    }
}

Complexity Analysis

Operation	Time Complexity	Space Complexity
Insert	O(L)	O(L)$$O(N \times L)
Search	O(L)	O(1)
Prefix Search	O(L)	O(1)

Where L is the length of the word and $N$ is the number of words.

Source Code

🚀 Code: All the code for this tutorial is available in my DSA-Kotlin Repository.