Contacts

  • + 0 comments
    class TrieNode:
        def __init__(self, data):
            self.data = data
            self.children = {}
            self.freq = 0
    
    def add_word(root, word):
        if len(word) == 0:
            return
        new_root = root.children.get(word[0], TrieNode(word[0]))
        new_root.freq = new_root.freq + 1
        root.children[word[0]] = new_root
        add_word(new_root, word[1:])
        
    def find_word(root, word):
        if len(word) == 0:
            return root.freq
        
        if word[0] not in root.children:
            return 0
            
        return find_word(root.children[word[0]], word[1:])
    
    def contacts(queries):
        found = []
        root = TrieNode('*')
        for query in queries:
            qtype, word = query[0], query[1]
            if qtype == 'add':
                add_word(root, word)
            else:
                found.append(find_word(root, word))
        return found