Leetcode 1804 Implement Trie II (Prefix Tree)

Problem statement

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

Solution

We need to keep TrieNode node with fields: children are another nodes, cnt_end is number of nodes end with this node and cnt_all is number of nodes go through this node.

1. update(word, h) will be used in insert and erase functions: for all nodes we update cnt_all and for the last we update cnt_end.
2. helper(word) will be used for countWordsEqualTo and countWordsStartingWith functions, where we traverse node by node and if it is not possible, return 0.

Complexity

Time complexity of all operations is O(n), where n is length of word. Space complexity is sum of lengths of all words.

Code

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

class Trie:
    def __init__(self):
        self.root = TrieNode()
        
    def update(self, word, h):
        node = self.root
        for symbol in word:
            node = node.children.setdefault(symbol, TrieNode())
            node.cnt_all += h
        node.cnt_end += h
        
    def helper(self, word):
        node = self.root
        for symbol in word:
            if symbol in node.children:
                node = node.children[symbol]
            else:
                return None
        return node
        
    def insert(self, word):
        self.update(word, 1)
        
    def erase(self, word):
        self.update(word, -1)
        
    def countWordsEqualTo(self, word):
        node = self.helper(word)
        return node.cnt_end if node != None else 0

    def countWordsStartingWith(self, word):
        node = self.helper(word)
        return node.cnt_all if node != None else 0