implement-trie-prefix-tree.py
trie/implement-trie-prefix-tree.py · Python · 1.9 KiB · 2026-02-17 21:07
# Used in prefix matching, autocomplete, spellchecker, etc. # tries # Could use hash map if the charset is not fixed. # O(n) time for insert and lookup, where n is the length of the word # O(t) space for total number of nodes created in a trie class Node: def __init__(self, c: str): self.val = c self.children = [None for _ in range(26)] self.is_end = False def contains_key(self, c: str) -> bool: return self.children[ord(c) - ord('a')] is not None def get_child(self, c: str) -> 'Node': return self.children[ord(c) - ord('a')] def set_child(self, c: str) -> None: self.children[ord(c) - ord('a')] = Node(c) def get_isend(self) -> bool: return self.is_end def set_isend(self) -> None: self.is_end = True class Trie: def __init__(self): self.root = Node('#') def insert(self, word: str) -> None: node = self.root for c in word: if not node.contains_key(c): node.set_child(c) node = node.get_child(c) node.set_isend() def search(self, word: str) -> bool: node = self.root for c in word: if not node.contains_key(c): return False node = node.get_child(c) return node.get_isend() def startswith(self, word: str) -> bool: node = self.root for c in word: if not node.contains_key(c): return False node = node.get_child(c) return True from unittest import TestCase class TrieTest(TestCase): def test_trie(self): t = Trie() t.insert("avocado") t.insert("telephone") self.assertTrue(t.search("avocado")) self.assertFalse(t.search("not-existend")) self.assertTrue(t.startswith("avo")) self.assertFalse(t.startswith("boh")) if __name__ == '__main__': import unittest unittest.main()