212. Word Search II

• Time: $O(mn \max(|\texttt{words[i]}|))$
• Space: $O(\Sigma |\texttt{words[i]}|)$

C++JavaPython

struct TrieNode {
  vector<shared_ptr<TrieNode>> children;
  const string* word = nullptr;
  TrieNode() : children(26) {}
};

class Solution {
 public:
  vector<string> findWords(vector<vector<char>>& board, vector<string>& words) {
    vector<string> ans;

    for (const string& word : words)
      insert(word);

    for (int i = 0; i < board.size(); ++i)
      for (int j = 0; j < board[0].size(); ++j)
        dfs(board, i, j, root, ans);

    return ans;
  }

 private:
  shared_ptr<TrieNode> root = make_shared<TrieNode>();

  void insert(const string& word) {
    shared_ptr<TrieNode> node = root;
    for (const char c : word) {
      const int i = c - 'a';
      if (node->children[i] == nullptr)
        node->children[i] = make_shared<TrieNode>();
      node = node->children[i];
    }
    node->word = &word;
  }

  void dfs(vector<vector<char>>& board, int i, int j, shared_ptr<TrieNode> node,
           vector<string>& ans) {
    if (i < 0 || i == board.size() || j < 0 || j == board[0].size())
      return;
    if (board[i][j] == '*')
      return;

    const char c = board[i][j];
    shared_ptr<TrieNode> child = node->children[c - 'a'];
    if (child == nullptr)
      return;
    if (child->word != nullptr) {
      ans.push_back(*child->word);
      child->word = nullptr;
    }

    board[i][j] = '*';
    dfs(board, i + 1, j, child, ans);
    dfs(board, i - 1, j, child, ans);
    dfs(board, i, j + 1, child, ans);
    dfs(board, i, j - 1, child, ans);
    board[i][j] = c;
  }
};

class TrieNode {
  public TrieNode[] children = new TrieNode[26];
  public String word;
}

class Solution {
  public List<String> findWords(char[][] board, String[] words) {
    for (final String word : words)
      insert(word);

    List<String> ans = new ArrayList<>();

    for (int i = 0; i < board.length; ++i)
      for (int j = 0; j < board[0].length; ++j)
        dfs(board, i, j, root, ans);

    return ans;
  }

  private TrieNode root = new TrieNode();

  private void insert(final String word) {
    TrieNode node = root;
    for (final char c : word.toCharArray()) {
      final int i = c - 'a';
      if (node.children[i] == null)
        node.children[i] = new TrieNode();
      node = node.children[i];
    }
    node.word = word;
  }

  private void dfs(char[][] board, int i, int j, TrieNode node, List<String> ans) {
    if (i < 0 || i == board.length || j < 0 || j == board[0].length)
      return;
    if (board[i][j] == '*')
      return;

    final char c = board[i][j];
    TrieNode child = node.children[c - 'a'];
    if (child == null)
      return;
    if (child.word != null) {
      ans.add(child.word);
      child.word = null;
    }

    board[i][j] = '*';
    dfs(board, i + 1, j, child, ans);
    dfs(board, i - 1, j, child, ans);
    dfs(board, i, j + 1, child, ans);
    dfs(board, i, j - 1, child, ans);
    board[i][j] = c;
  }
}

class TrieNode:
  def __init__(self):
    self.children: dict[str, TrieNode] = {}
    self.word: str | None = None


class Solution:
  def findWords(self, board: list[list[str]], words: list[str]) -> list[str]:
    m = len(board)
    n = len(board[0])
    ans = []
    root = TrieNode()

    def insert(word: str) -> None:
      node = root
      for c in word:
        node = node.children.setdefault(c, TrieNode())
      node.word = word

    for word in words:
      insert(word)

    def dfs(i: int, j: int, node: TrieNode) -> None:
      if i < 0 or i == m or j < 0 or j == n:
        return
      if board[i][j] == '*':
        return

      c = board[i][j]
      if c not in node.children:
        return

      child = node.children[c]
      if child.word:
        ans.append(child.word)
        child.word = None

      board[i][j] = '*'
      dfs(i + 1, j, child)
      dfs(i - 1, j, child)
      dfs(i, j + 1, child)
      dfs(i, j - 1, child)
      board[i][j] = c

    for i in range(m):
      for j in range(n):
        dfs(i, j, root)

    return ans