Read-only web browser

Solutions

Browse the interview solutions folder with syntax highlighting.

implement-trie-prefix-tree.py

trie/implement-trie-prefix-tree.py · Python · 1.9 KiB · 2026-02-17 21:07

Back to folder
# 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()