Skip to content

A prefix tree (also known as a trie) is a tree data structure used to efficiently store and retrieve keys in a set of strings. Some applications of this data structure include auto-complete and spell checker systems.

Implement the PrefixTree class:

  • PrefixTree() Initializes the prefix tree object.
  • void insert(String word) Inserts the string word into the prefix tree.
  • boolean search(String word) Returns true if the string word is in the prefix tree (i.e., was inserted before), and false otherwise.
  • boolean startsWith(String prefix) Returns true if there is a previously inserted string word that has the prefix prefix, and false otherwise.

Trie Template (lowercase english characters)

The Nodes are slotted into right position of the children of its parent node, which represents the previous character.

Note: all ending node is going to have an empty children array, and its end flag is set to True.

python
class PrefixTreeNode:
    def __init__(self):
        self.children = {}
        self.end = False

class Trie:

    def __init__(self):
        self.root = PrefixTreeNode()

    def insert(self, word: str) -> None:
        curr = self.root
        for c in word:
            if c not in curr.children:
                curr.children[c] = PrefixTreeNode()
            curr = curr.children[c]
        curr.end = True

    def search(self, word: str) -> bool:
        curr = self.root
        for c in word:
            if c not in curr.children:
                return False
            curr = curr.children[c]
        return curr.end

    def startsWith(self, prefix: str) -> bool:
        curr = self.root
        for c in prefix:
            if c not in curr.children:
                return False
            curr = curr.children[c]
        return True
python
class Node:
    def __init__(self):
        self.children = [None] * 26
        self.end = False

class WordDictionary:

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

    def addWord(self, word: str) -> None:
        curr = self.root
        for c in word:
            i = ord(c) - ord('a')
            if not curr.children[i]:
                curr.children[i] = Node()
            curr = curr.children[i]
        curr.end = True

    def searchAt(self, startNode, word) -> bool:
        if not word:
            return startNode.end
        topc = word[0]
        if topc == '.':
            for child in startNode.children:
                if child and self.searchAt(child, word[1:]):
                    return True
        else:
            idx = ord(topc) - ord('a')
            if not startNode.children[idx]:
                return False
            return self.searchAt(startNode.children[idx], word[1:])
        return False

    def search(self, word: str) -> bool:
        if word == "":
            return self.root.end
        return self.searchAt(self.root, word)

Word search ii

python
class TrieNode:
    def __init__(self):
        # Each TrieNode has a dictionary to store its children and a boolean to indicate if it's the end of a word
        self.children = {}
        self.end = False

    def addWord(self, word):
        # Adds a word to the Trie by iterating through each character
        cur = self
        for c in word:
            # If the character is not a child of the current node, create a new TrieNode for it
            if c not in cur.children:
                cur.children[c] = TrieNode()
            # Move to the child node
            cur = cur.children[c]
        # Mark the last node as the end of a word
        cur.end = True

# Directions for moving in the board (up, down, left, right)
DIR = {(0, -1), (0, 1), (1, 0), (-1, 0)}

class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        # Initialize the root of the Trie and add all words to it
        root = TrieNode()
        for w in words:
            root.addWord(w)

        # Get the dimensions of the board
        m, n = len(board), len(board[0])
        # Set to store the result words and to keep track of visited cells
        res, visited = set(), set()

        def dfs(r, c, node, word):
            # If the current cell's character is not in the current node's children or the cell is visited, return
            if board[r][c] not in node.children or (r, c) in visited:
                return

            # Mark the cell as visited
            visited.add((r, c))
            # Move to the child node corresponding to the current cell's character
            node = node.children[board[r][c]]
            # Append the character to the current word
            word += board[r][c]
            # If the current node marks the end of a word, add it to the result set
            if node.end:
                res.add(word)

            # Recursively call dfs for all valid neighboring cells
            for (dx, dy) in DIR:
                nx, ny = r+dx, c+dy
                if 0 <= nx < m and 0 <= ny < n:
                    dfs(nx, ny, node, word)
            # Backtrack by removing the current cell from the visited set
            visited.remove((r, c))

        # Iterate over all cells in the board and call dfs starting from each cell
        for r in range(m):
            for c in range(n):
                dfs(r, c, root, "")

        # Convert the result set to a list and return it
        return list(res)