Tries: Contacts

  • + 0 comments

    Trie solution in python - storing number of matches in each node when inserting

    class TrieNode:
        def __init__(self, char):
            self.char = char
            self.children = {}
            self.valid_children = 0
    
    
    class Trie:
        def __init__(self):
            self.root = TrieNode("")
        
        def insert(self, node, word):
            if not word:
                new_node = TrieNode("*")
                node.children["*"] = new_node
                node.valid_children += 1
                return 1
            elif word[0] in node.children:
                new_node = node.children[word[0]]
            else:
                new_node = TrieNode(word[0])
                node.children[word[0]] = new_node
            valid_children = self.insert(new_node, word[1:])
            node.valid_children += valid_children
            return valid_children
        
        def find_matches(self, query):
            node = self.root
            
            for char in query:
                if char in node.children:
                    node = node.children[char]
                else:
                    return 0
            return node.valid_children 
        
        def number_of_valid_children(self, prefix):
            return self.find_matches(prefix)
    
    def contacts(queries):
        trie = Trie()
        result = []
        for query in queries:
            op, value = query[0], query[1]
            if op == "add":
                trie.insert(trie.root, value)
            elif op == "find":
                result.append(trie.number_of_valid_children(value))
        return result