208. Implement Trie (Prefix Tree) - LeetCode

from typing import Optional


class Node:
    def __init__(self, is_end: bool = False):
        self.is_end = is_end
        self.children: dict[str, Node] = {}

    def __repr__(self):
        return f'''{'! ' if self.is_end else ''}{self.children}'''


class Trie:
    def __init__(self):
        self.root = Node()

    def insert(self, word: str) -> None:
        node = self.root
        for c in word:
            node.children[c] = node.children.get(c, Node())
            node = node.children[c]
        node.is_end = True

    def _last_node_same_with(self, word: str) -> Optional[Node]:
        node = self.root
        for c in word:
            if c not in node.children:
                return
            node = node.children[c]
        return node

    def search(self, word: str) -> bool:
        node = self._last_node_same_with(word)
        return bool(node) and node.is_end

    def startsWith(self, word: str) -> bool:
        node = self._last_node_same_with(word)
        return bool(node)