Leetcode 1397 Solution

class TrieNode:
    def __init__(self):
        self.children = {}
        self.words = []

    def append_word(self, word):
        self.words.append(word)
        self.words = sorted(self.words)[:3]

    def get_child(self, ch):
        if ch not in self.children:
            self.children[ch] = TrieNode()
        return self.children[ch]


class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        root = TrieNode()

        for product in products:
            curr_node = root
            for ch in product:
                curr_node = curr_node.get_child(ch)
                curr_node.append_word(product)

        curr_node = root
        ans = []
        for ch in searchWord:
            curr_node = curr_node.get_child(ch)
            ans.append(curr_node.words)

        return ans