|
| 1 | +""" |
| 2 | +A Trie/Prefix Tree is a kind of search tree used to provide quick lookup |
| 3 | +of words/patterns in a set of words. A basic Trie however has O(n^2) space complexity |
| 4 | +making it impractical in practice. It however provides O(max(search_string, length of longest word)) lookup |
| 5 | +time making it an optimal approach when space is not an issue. |
| 6 | +
|
| 7 | +""" |
| 8 | + |
| 9 | + |
| 10 | +class TrieNode: |
| 11 | + def __init__(self): |
| 12 | + self.nodes = dict() # Mapping from char to TrieNode |
| 13 | + self.is_leaf = False |
| 14 | + |
| 15 | + def insert_many(self, words: [str]): |
| 16 | + """ |
| 17 | + Inserts a list of words into the Trie |
| 18 | + :param words: list of string words |
| 19 | + :return: None |
| 20 | + """ |
| 21 | + for word in words: |
| 22 | + self.insert(word) |
| 23 | + |
| 24 | + def insert(self, word: str): |
| 25 | + """ |
| 26 | + Inserts a word into the Trie |
| 27 | + :param word: word to be inserted |
| 28 | + :return: None |
| 29 | + """ |
| 30 | + curr = self |
| 31 | + for char in word: |
| 32 | + if char not in curr.nodes: |
| 33 | + curr.nodes[char] = TrieNode() |
| 34 | + curr = curr.nodes[char] |
| 35 | + curr.is_leaf = True |
| 36 | + |
| 37 | + def find(self, word: str) -> bool: |
| 38 | + """ |
| 39 | + Tries to find word in a Trie |
| 40 | + :param word: word to look for |
| 41 | + :return: Returns True if word is found, False otherwise |
| 42 | + """ |
| 43 | + curr = self |
| 44 | + for char in word: |
| 45 | + if char not in curr.nodes: |
| 46 | + return False |
| 47 | + curr = curr.nodes[char] |
| 48 | + return curr.is_leaf |
| 49 | + |
| 50 | + |
| 51 | +def print_words(node: TrieNode, word: str): |
| 52 | + """ |
| 53 | + Prints all the words in a Trie |
| 54 | + :param node: root node of Trie |
| 55 | + :param word: Word variable should be empty at start |
| 56 | + :return: None |
| 57 | + """ |
| 58 | + if node.is_leaf: |
| 59 | + print(word, end=' ') |
| 60 | + |
| 61 | + for key, value in node.nodes.items(): |
| 62 | + print_words(value, word + key) |
| 63 | + |
| 64 | + |
| 65 | +def test(): |
| 66 | + words = ['banana', 'bananas', 'bandana', 'band', 'apple', 'all', 'beast'] |
| 67 | + root = TrieNode() |
| 68 | + root.insert_many(words) |
| 69 | + # print_words(root, '') |
| 70 | + assert root.find('banana') |
| 71 | + assert not root.find('bandanas') |
| 72 | + assert not root.find('apps') |
| 73 | + assert root.find('apple') |
| 74 | + |
| 75 | +test() |
0 commit comments