Tries: Contacts

  • + 0 comments

    Working Python recursive OOP implementation of Trie. Storing the size rather than computing recursively is essential to pass the tests.

    class PrefixTree():
        
        def __init__(self, is_end=False):
            self.d = dict()
            self.is_end = is_end
            self.size = 0
        
        # Recursively add character by character
        def add(self, word):
            # Base case: empty or existing word
            if len(word) == 0 or self.find(word):
                return -1
            if word[0] not in self.d.keys():
                self.d[word[0]] = PrefixTree()
            # Base case: one character word
            if len(word) == 1:
                self.d[word[0]].is_end = True
            # Recursive case: add the remaining characters
            else:
                self.d[word[0]].add(word[1:])
            # Caching the size is the key to good runtime in 'find()'
            self.size += 1
        
        # Returns number of matches for a prefix
        def find(self, prefix):
            # Base case: we're at the end node of the prefix
            if len(prefix) == 0:
                return self.size + (1 if self.is_end else 0)
            # Base case: we can't reach the next character
            if prefix[0] not in self.d.keys():
                return 0
            # Recursive case: keep looking
            return self.d[prefix[0]].find(prefix[1:])