Source code for pygorithm.data_structures.trie

"""
Author: MrDupin
"""

class Node:
    def __init__(self, v, p=None, w=False):
        self.word = w #If the node represents the end of a word or not
        self.parent = p
        self.value = v
        self.children = {}


[docs]class Trie: def __init__(self): self.root = Node('') #The root of the trie is always empty
[docs] def insert(self, word): """ Inserts a word in the trie. Starting from the root, move down the trie following the path of characters in the word. If the nodes for the word characters end, add them. When the last char is added, mark it as a word-ending node. """ l = len(word) curr = self.root for i, c in enumerate(word): last = False if(i == l-1): #The last char of the word last = True if(c not in curr.children): curr.children[c] = Node(c, curr, last) elif(last): #c already exists, but as it is the last char of word, #it should now be flagged as a word in the trie. curr.children[c].word = True curr = curr.children[c]
[docs] def search(self, word): """ Searches for given word in trie. We want to find the last node for the word. If we can't, then it means the word is not in the trie. """ if self.find_final_node(word): return True else: return False
[docs] def find_words(self, prefix): """ Find all words with the given prefix """ v = self.find_final_node(prefix) wList = self.build_word_list(v, prefix) if(v and v.word): #v exists and the prefix is itself a word; add it to the list. wList.append(prefix) return wList
[docs] def find_final_node(self, word): """ Returns the last node in given word. The process goes like this: Start from the root. For every char in word, go down one level. If we can't go down a level, then the word doesn't exist. If we do, and the current char is the last char of the word and the node we are currently at is a word, then we have found the given word. """ curr = self.root l = len(word) for i, c in enumerate(word): if(c not in curr.children): #There is no prefix of cWord + c return None if(i == l-1): #Last char of word return curr.children[c] curr = curr.children[c] return None
[docs] def build_word_list(self, v, cWord): """ Recursively builds the list of words. * v: Node to check * cWord : The word built up to v """ if(not v): return None wList = [] for i, k in v.children.items(): tempWord = cWord + i if(k.word): #If the iterated prefix is a word, add it to the list wList.append(tempWord) #The list of words under tWord wList.extend(self.build_word_list(k, tempWord)) return wList