P145 实现前缀树

描述#

这个题目说的是，你要实现一棵前缀树（Trie），其中包含 insert, search 和 startsWith 三个方法。insert 用于插入一个单词，search 用于搜索一个单词是否存在于前缀树中，startsWith 用于判断前缀树中是否存在一个包含给定前缀的单词。注意：所有的输入都是非空字符串，并且只包含小写字母。

代码#

public class AlgoCasts {

  public class Trie {

    private class TrieNode {
      boolean endOfWord;
      TrieNode[] children;
      TrieNode() {
        this.endOfWord = false;
        this.children = new TrieNode[26];
      }
    }

    private TrieNode root;

    /** Initialize your data structure here. */
    public Trie() {
      this.root = new TrieNode();
    }

    /** Inserts a word into the trie. */
    public void insert(String word) {
      TrieNode cur = root;
      for (int i = 0; i < word.length(); ++i) {
        int idx = word.charAt(i) - 'a';
        if (cur.children[idx] == null) cur.children[idx] = new TrieNode();
        cur = cur.children[idx];
      }
      cur.endOfWord = true;
    }

    /** Returns if the word is in the trie. */
    public boolean search(String word) {
      TrieNode node = searchEndNode(word);
      return node != null && node.endOfWord;
    }

    /** Returns if there is any word in the trie that starts with the given prefix. */
    public boolean startsWith(String prefix) {
      TrieNode node = searchEndNode(prefix);
      return node != null;
    }

    private TrieNode searchEndNode(String str) {
      TrieNode cur = root;
      for (int i = 0; i < str.length() && cur != null; ++i) {
        int idx = str.charAt(i) - 'a';
        cur = cur.children[idx];
      }
      return cur;
    }

  }

}