前缀树（trie）实现

简介

前缀树（trie）又称字典树、单词查找树。它由字符串键中的所有字符构造而成，允许使用被查找键中的字符进行查找。它的英文单词 trie 来自于 E.Fredkin 在 1960 年玩的一个文字游戏，因为这个数据结构的作用是取出（ retrieval）数据，但发音为 try 是为了避免与 tree 相混淆。

性质

和其它查找树一样，单词查找树也是由链接的结点所组成的数据结构，这些链接可能为空，也可能指向其他结点。每个结点都只可能有一个指向它的结点，称为它的父结点（只有一个结点除外，即根结点，没有任何结点指向根结点）。每个结点都含有 R 条链接，其中 R 为字母表的大小。 所以前缀树其本质就是R叉树。

算法思路

• 假设字母表大小只有26个小写英文字母，指向子节点的指针数组children的长度为 26，即小写英文字母的数量。此时children[0] 对应小写字母 a，children[1]对应小写字母 b，…，children[25] 对应小写字母 z。
• 此外还需要一个标志来判断是否为单词的结尾。考虑上图中she和shell单词在e的节点有个标志来表示为she单词的结尾，在最后一个字符l有个标志来表示为shell单词的结尾。

插入字符串

我们从前缀树的根开始，插入字符串。对于当前字符对应的子节点，有两种情况：

• 子节点存在。沿着指针移动到子节点，继续处理下一个字符。

• 子节点不存在。创建一个新的子节点，记录在 children 数组的对应位置上，然后沿着指针移动到子节点，继续搜索下一个字符。

重复以上步骤，直到处理字符串的最后一个字符，然后将当前节点标记为字符串的结尾。

查找前缀

我们从字典树的根开始，查找前缀。对于当前字符对应的子节点，有两种情况：

• 子节点存在。沿着指针移动到子节点，继续搜索下一个字符。

• 子节点不存在。说明字典树中不包含该前缀，返回空指针。

重复以上步骤，直到返回空指针或搜索完前缀的最后一个字符。若搜索到了前缀的末尾，就说明字典树中存在该前缀。此外，若前缀末尾对应节点的 isEnd 为真，则说明字典树中存在该字符串。

代码实现

实现类

package com.design.trie;

/**
 * 前缀树，本质上是多叉树（26叉树）
 *
 *  @author hh
 *  @date 2021-12-30 12:16
 */
public class Trie {

    /**
     * 根节点
     */
    private TrieNode root;

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

    /**
     * 插入单词
     *
     * @param word 单词
     */
    public void insert(String word){
        // TODO: 2021-12-30 校验单词是否为空，暂时省略
        TrieNode p = this.root;
        for(char c : word.toCharArray()){
            //如果没有构建节点，则构建节点
            if(p.getChildren()[c - 'a'] == null){
                p.getChildren()[c - 'a'] = new TrieNode();
            }
            p = p.getChildren()[c - 'a'];
        }
        p.setEnd(true);
    }

    /**
     * 查找单词，返回是否存在
     *
     * @param word 单词
     * @return true或false
     */
    public boolean search(String word){
        //结束条件是到达叶子节点，如果没有到达叶子节点也返回false
        TrieNode p = this.match(word);
        return  p != null && p.isEnd();
    }

    /**
     * 前缀树是否存在以prefix前缀开始的单词
     *
     * @param prefix 前缀
     * @return true或false
     */
    public boolean prefixStartWith(String prefix){
        //结束条件是到达叶子节点，如果没有到达叶子节点也返回false
        return this.match(prefix) != null;
    }

    /**
     * 检查前缀树中是否匹配text文本
     *
     * @param text 文本串
     * @return true或false
     */
    private TrieNode match(String text){
        TrieNode p = this.root;
        for(char c : text.toCharArray()){
            int index = c - 'a';
            if(p.getChildren()[index] != null){
                p = p.getChildren()[index];
            }else{
                return null;
            }
        }
        return p;
    }


    static class TrieNode{
        /**
         * 孩子节点数组
         */
        private TrieNode[] children;

        /**
         * 是否为叶子节点
         */
        private boolean isEnd;

        public TrieNode() {
            this.children  = new TrieNode[26];
            this.isEnd = false;
        }

        public TrieNode[] getChildren() {
            return children;
        }

        public void setChildren(TrieNode[] children) {
            this.children = children;
        }

        public boolean isEnd() {
            return isEnd;
        }

        public void setEnd(boolean end) {
            isEnd = end;
        }
    }

}

测试类

package com.design.trie;

public class TrieTest {

    public static void main(String[] args){
        Trie trie = new Trie();
        trie.insert("love");
        trie.insert("lovey");
        System.out.println(trie.prefixStartWith("lov"));
        System.out.println(trie.search("love"));
        System.out.println(trie.search("lovey"));
        System.out.println(trie.search("ylove"));
    }
}

测试结果

参考文档

• 《算法(第4版)》5.2节 单词查找树
• leetcode