面试算法62：实现前缀树

题目

请设计实现一棵前缀树Trie，它有如下操作。

• 函数insert，在前缀树中添加一个字符串。
• 函数search，查找字符串。如果前缀树中包含该字符串，则返回true；否则返回false。
• 函数startWith，查找字符串前缀。如果前缀树中包含以该前缀开头的字符串，则返回true；否则返回false。

例如，调用函数insert在前缀树中添加单词"goodbye"之后，输入"good"调用函数search返回false，但输入"good"调用函数startWith则返回true。再次调用函数insert添加单词"good"之后，此时再输入"good"调用函数search则返回true。

分析

首先定义前缀树中节点的数据结构。前缀树中的节点对应字符串中的一个字符。如果只考虑英文小写字母，那么字符可能是从’a’到’z’的任意一个，因此前缀树中的节点可能有26个子节点。可以将26个子节点放到一个数组中，数组中的第1个元素是对应字母’a’的子节点，第2个元素是对应字母’b’的子节点，其余的以此类推。

解

public class Trie {

    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    class TrieNode {
        TrieNode children[];
        boolean isWord;

        public TrieNode() {
            children = new TrieNode[26];
        }
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("apple");
        System.out.println(trie.search("apple"));
        System.out.println(trie.search("app"));
        System.out.println(trie.startsWith("app"));
        trie.insert("app");
        System.out.println(trie.search("app"));
    }

    public void insert(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            // 对应的位置存放一个链接，相应位置有值则代表有相应的字母
            if (node.children[ch - 'a'] == null) {
                node.children[ch - 'a'] = new TrieNode();
            }

            node = node.children[ch - 'a'];
        }

        node.isWord = true;
    }

    public boolean search(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            if (node.children[ch - 'a'] == null) {
                return false;
            }

            node = node.children[ch - 'a'];
        }

        return node.isWord;
    }

    public boolean startsWith(String prefix) {
        TrieNode node = root;
        for (char ch : prefix.toCharArray()) {
            if (node.children[ch - 'a'] == null) {
                return false;
            }

            node = node.children[ch - 'a'];
        }

        return true;
    }
}