Prefix Tree (Trie) [LeetCode 208] [LeetCode 1268]

Trie.py

from typing import List

class TrieNode:
    def __init__(self):
        self.next = dict() # pointer to next node, starting with a key char
        self.end = False # by default false
        
class Trie:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.root = TrieNode() # Make a TrieNode
        return

    def insert(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        curr = self.root # start with root
        for char in word:
            # If char is not in next, add a new node to next
            if char not in curr.next:
                curr.next[char] = TrieNode()
            curr = curr.next[char] # point to next node
        curr.end = True # Mark end of word
        return
        

    def search(self, word: str) -> bool:
        """
        Returns if the word is in the trie.
        """
        curr = self.root # start with root
        for char in word:
            # If char not in next, return false
            if char not in curr.next:
                return False
            curr = curr.next[char] # point to next node
        return curr.end # return end of word
            
    def searchAbsolute(self, word: str, curr: TrieNode = None) -> bool:
        """
        Returns if the word is in the trie, with wild characters like "."
        """

        if not curr:
            curr = self.root # start with root

        for i,char in enumerate(word):

            # char is ".", search all nodes under the current node
            if char == ".":
                for next in curr.next.values():
                    if self.searchAbsolute(word[i+1:], next):
                        return True
                return False # if word not found, return false

            # If char in next, move to next
            elif char in curr.next:
                curr = curr.next[char] # point to next node

            # If char not in next, return false
            else:
                return False

        return curr.end # return end of word

    def startsWith(self, prefix: str) -> bool:
        """
        Returns if there is any word in the trie that starts with the given prefix.
        """
        curr = self.root # start with root
        for char in prefix:
            # If char not in next, return false
            if char not in curr.next:
                return False
            curr = curr.next[char] # point to next node
        return True # by default return true
    
    def searchWithPrefix(self, prefix: str, k: int = -1) -> List:
        """
        For k > 0 returns first k words starting with prefix.
        For k == -1 returns all words starting with prefix.
        By default k == -1
        """

        words = []
        
        # Search and save all words starting with prefix
        def searchWords(curr, prefix):
            
            nonlocal words
            nonlocal k
            
            if k == 0: # Terminate condition for k > 0
                return

            if curr.end:
                words.append(prefix)
                k = k - 1

            # Search for next word
            for char in sorted(curr.next.keys()): # maintain lexicographical order
                searchWords(curr.next[char], prefix + char)
                
            return # ~searchWords()

        # Search if prefix exist in tree
        curr = self.root
        for char in prefix:
            if char in curr.next:
                curr = curr.next[char]
            else:
                return words

        # Search from current point till end of each word which starts with prefix
        searchWords(curr, prefix)

        return words