208. Implement Trie (Prefix Tree)

Implement a trie with insert, search, and startsWith methods.

Note:
You may assume that all inputs are consist of lowercase letters a-z.

class TrieNode {
  public char val;
  public boolean isWord;
  public TrieNode[] children;
  public TrieNode(char c) {
    this.val = c;
    this.isWord = false;
    this.children = new TrieNode[26];
  }
}

public class Trie {
  private TrieNode root;
  public Trie() {
    root = new TrieNode(' ');
  }
  
  public void insert(String word) {
    TrieNode tn = root;
    for (int i = 0; i < word.length(); i++) {
      char c = word.charAt(i);
      if (tn.children[c - 'a'] == null)
        tn.children[c - 'a'] = new TrieNode(c);
      tn = tn.children[c - 'a'];
    }
    tn.isWord = true;
  }
  
  public boolean search(String word) {
    TrieNode tn = root;
    for (int i = 0; i < word.length(); i++) {
      char c = word.charAt(i);
      if (tn.children[c - 'a'] == null)
        return false;
      tn = tn.children[c - 'a'];
    }
    return tn.isWord;
  }
  
  public boolean startsWith(String prefix) {
    TrieNode tn = root;
    for(int i = 0; i < prefix.length(); i++) {
      char c = prefix.charAt(i);
      if (tn.children[c - 'a'] == null)
        return false;
      tn = tn.children[c - 'a'];
    }
    return true;
  }
}

public List<String> match(String prefix) {
    TrieNode tn = root;
    // 第一步：遍历前缀，找到前缀最后一个字符对应的节点
    for (int i = 0; i < prefix.length(); i++) {
        char ch = prefix.charAt(i);
        // 如果前缀中的字符不存在，直接返回空列表
        if (tn.children[ch - 'a'] == null) {
            return new ArrayList<>();
        }
        tn = tn.children[ch - 'a'];
    }

    // 第二步：从前缀末尾节点开始，递归收集所有完整单词
    List<String> res = new ArrayList<>();
    helper(tn, new StringBuilder(), res, prefix);
    return res;
}

// 递归辅助方法：深度优先遍历子节点，收集所有完整单词
private void helper(TrieNode node, StringBuilder path, List<String> res, String prefix) {
    // 如果当前节点是一个单词的结尾，拼接前缀和路径，加入结果集
    if (node.isWord) {
        res.add(prefix + path.toString());
    }

    for (TrieNode child : node.children) {
        if (child == null) {
            continue;
        }
        path.append(child.val);
        helper(child, path, res, prefix);
        // 回溯：删除最后一个字符
        path.deleteCharAt(path.length() - 1);
    }
}