Java 前缀树Trie的实现

Trie的基础实现：本实现只适用于英语这么语言，将单词存放在Character对象中，所以Trie才没有使用泛型。

使用泛型是比较好的设计：

不过这里的应用场景为英语单词，所以这里不采用泛型。这里将单词看成字符串是由一个个Character组成的。

1.向Trie添加元素（非递归写法）

import java.util.TreeMap;

public class Trie {

    private class Node{

        public boolean isWord;
        //到下一个节点的映射
        public TreeMap<Character, Node> next;

        public Node(boolean isWord){
            this.isWord = isWord;
            //初始化字典树
            next = new TreeMap<>();
        }

        public Node(){
            this(false);
        }
    }
    //根节点
    private Node root;
    //Trie单词个数
    private int size;

    public Trie(){
        root = new Node();
        size = 0;
    }

    // 获得Trie中存储的单词数量
    public int getSize(){
        return size;
    }

    // 向Trie中添加一个新的单词word
    //将单词拆分成一个个字符c，然后从根节点开始往下添加
    public void add(String word){

        Node cur = root;
        //循环判断新的cur节点是否包含下一个字符到下一个节点的映射
        for(int i = 0 ; i < word.length() ; i ++){
            //将c当成一个节点插入Trie中
            char c = word.charAt(i);
            //判断cur.next是不是已经指向我们要找的c字符相应的节点
            if(cur.next.get(c) == null){
                //新建节点
                cur.next.put(c, new Node());
            }
            //否则，就直接走到该节点位置即可
            cur = cur.next.get(c);
        }
        //判断该单词并不表示任何一个单词的结尾
        if(!cur.isWord){
            //确定cur是新的单词
            cur.isWord = true;
            size ++;
        }
    }
}

2.从Trie中查询元素

    // 查询单词word是否在Trie中
    public boolean contains(String word){

        Node cur = root;
        //从单词的首字母的下一个节点出发，一直按顺序的寻找该单词的字符
        //如果该字符的节点不存在，则返回false
        //直到叶子节点位置
        for(int i = 0 ; i < word.length() ; i ++){
            char c = word.charAt(i);
            if(cur.next.get(c) == null){
                return false;
            }
            //进行下一个节点的查找
            cur = cur.next.get(c);
        }
        //返回是否找到word
        return cur.isWord;
    }

Trie只能存储字符串型元素集合。而之前学习的基于二分搜索树实现的集合可以存储任意类型的元素，所以下面对这两种数据结构做一下性能的比较：这里对《傲慢与偏见》这本书进行单词的查找。

文件读取类：

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Locale;
import java.util.Scanner;

// 文件相关操作
public class FileOperation {

    // 读取文件名称为filename中的内容，并将其中包含的所有词语放进words中
    public static boolean readFile(String filename, ArrayList<String> words){

        if (filename == null || words == null){
            System.out.println("filename is null or words is null");
            return false;
        }

        // 文件读取
        Scanner scanner;

        try {
            File file = new File(filename);
            if(file.exists()){
                FileInputStream fis = new FileInputStream(file);
                scanner = new Scanner(new BufferedInputStream(fis), "UTF-8");
                scanner.useLocale(Locale.ENGLISH);
            }else{
                return false;
            }
        } catch(IOException ioe){
            System.out.println("Cannot open " + filename);
            return false;
        }

        // 简单分词
        // 这个分词方式相对简陋, 没有考虑很多文本处理中的特殊问题
        // 在这里只做demo展示用
        if (scanner.hasNextLine()) {

            String contents = scanner.useDelimiter("\\A").next();

            int start = firstCharacterIndex(contents, 0);
            for (int i = start + 1 ; i <= contents.length() ; )
                if (i == contents.length() || !Character.isLetter(contents.charAt(i))) {
                    String word = contents.substring(start, i).toLowerCase();
                    words.add(word);
                    start = firstCharacterIndex(contents, i);
                    i = start + 1;
                }else{
                    i ++;
                }
        }

        return true;
    }

    // 寻找字符串s中，从start的位置开始的第一个字母字符的位置
    private static int firstCharacterIndex(String s, int start){

        for( int i = start ; i < s.length() ; i ++ ){
            if( Character.isLetter(s.charAt(i)) ){
                return i;
            }
        }
        return s.length();
    }
}

测试类：

import java.util.ArrayList;

public class Main {

    public static void main(String[] args) {

        System.out.println("Pride and Prejudice");
        
        ArrayList<String> words = new ArrayList<>();
        //将从文件读取类中读取到的单词往数组中添加
        if(FileOperation.readFile("pride-and-prejudice.txt", words)){

            long startTime = System.nanoTime();

            BSTSet<String> set = new BSTSet<>();
            for(String word: words){
                set.add(word);
            }
            //查找与一下添加进去的单词，为了计算出该操作的时间，方便比较
            for(String word: words){
                set.contains(word);
            }
            long endTime = System.nanoTime();

            double time = (endTime - startTime) / 1000000000.0;

            System.out.println("Total different words: " + set.getSize());
            System.out.println("BSTSet: " + time + " s");

            // ---

            startTime = System.nanoTime();

            Trie trie = new Trie();
            for(String word: words){
                trie.add(word);
            }
            for(String word: words){
                trie.contains(word);
            }
            endTime = System.nanoTime();

            time = (endTime - startTime) / 1000000000.0;

            System.out.println("Total different words: " + trie.getSize());
            System.out.println("Trie: " + time + " s");
        }
    }
}

测试结果：

Java 前缀树Trie的实现

结果发现Trie单词的查找能力要略高于BSTSet集合。

3.字典树Trie的前缀查询

Trie的搜索过程本身就是在查找单词前缀的过程。

    // 查询是否在Trie中有单词以prefix为前缀
    public boolean isPrefix(String prefix){

        Node cur = root;
        for(int i = 0 ; i < prefix.length() ; i ++){
            char c = prefix.charAt(i);
            if(cur.next.get(c) == null){
                return false;
            }
            cur = cur.next.get(c);
        }

        return true;
    }