Leetcode 0208 Implement Trie (Prefix Tree)

Problem statement

https://leetcode.com/problems/implement-trie-prefix-tree/

Solution

Define class TrieNode, where we have two fields: self.children = {} and self.endNode = 0. Then all three desired operations are straightforward: just go symbol by symbol and insert if needed or search.

Complexity

Time complexity of all operations is O(n), where n is the size of word. Space complexity is O(M), where M is total lengths of all words.

Code

class TrieNode:
    def __init__(self):
        self.children = {}
        self.end_node = 0

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

    def insert(self, word): #None
        root = self.root
        for symbol in word:
            root = root.children.setdefault(symbol, TrieNode())
        root.end_node = 1
        
    def searchHelper(self, word):
        root = self.root
        for symbol in word:
            if symbol in root.children:
                root = root.children[symbol]
            else:
                return -1

        return 1 if root.end_node == 1 else 0  

    def search(self, word): #bool:
        return self.searchHelper(word) == 1

    def startsWith(self, prefix): #bool:
        return self.searchHelper(prefix) >= 0