Implement Trie Prefix Tree

Question

trie (pronounced as "try") or prefix tree is a tree data structure used to efficiently store and retrieve keys in a dataset of strings. There are various applications of this data structure, such as autocomplete and spellchecker.

Implement the Trie class:

  • Trie() Initializes the trie object.
  • void insert(String word) Inserts the string word into the trie.
  • boolean search(String word) Returns true if the string word is in the trie (i.e., was inserted before), and false otherwise.
  • boolean startsWith(String prefix) Returns true if there is a previously inserted string word that has the prefix prefix, and false otherwise.

Example:

Input
["Trie", "insert", "search", "search", "startsWith", "insert", "search"]
[[], ["apple"], ["apple"], ["app"], ["app"], ["app"], ["app"]]
Output
[null, null, true, false, true, null, true]

Explanation
Trie trie = new Trie();
trie.insert("apple");
trie.search("apple");   // return True
trie.search("app");     // return False
trie.startsWith("app"); // return True
trie.insert("app");
trie.search("app");     // return True

Solution

from typing import *
class TrieNode:
    val: int
    neighbours: Dict[str, 'TrieNode']
    isWord: bool

    def __init__(self, val: str):
        self.val = val
        self.neighbours = {} #  { a: TrieNode(a) } 
        self.isWord = False

class Trie:
    def __init__(self):
        self.root = TrieNode("*")        

    def insert(self, word: str) -> None:
        currNode = self.root
        for letter in word:
            if letter not in currNode.neighbours:
                currNode.neighbours[letter] = TrieNode(letter)
            currNode = currNode.neighbours[letter]
        currNode.isWord = True

    def search(self, word: str) -> bool:
        currNode = self.root
        for letter in word:
            if letter not in currNode.neighbours:
                return False
            
            currNode = currNode.neighbours[letter]
        
        return currNode.isWord
        
    def startsWith(self, prefix: str) -> bool:
        currNode = self.root
        for letter in prefix:
            if letter not in currNode.neighbours:
                return False
            currNode = currNode.neighbours[letter]
        return True

Time complexity: $O(L)$ — $L$ is the length of the longest word

Space complexity: $O(L)$ — $L$ is the lenght of the longest word, $L$ is the length of the tree

def search(self, word: str) -> bool:
	return self.recursiveSearch(word, 0, self.root)

def recursiveSearch(self, word: str, index: int, node: TrieNode) -> bool:
	if index >= len(word): return False

	letter = word[index]
	if letter not in node.neighbours: return False

	letterNode = node.neighbours[letter]
	if index == len(word) - 1 and letterNode.isWord:
		return True
	
	return self.recursiveSearch(word, index + 1, letterNode)

Recursively startWith

def startsWith(self, prefix: str) -> bool:
	return self.recursiveStartWith(prefix, 0, self.root)

def recursiveStartWith(self, prefix, index, node) -> bool:
	if index >= len(prefix): return False

	letter = prefix[index]
	if letter not in node.neighbours: return False

	if index == len(prefix) - 1:
		return True
	
	letterNode = node.neighbours[letter]
	return self.recursiveStartWith(prefix, index + 1, letterNode)