java 实现字典树（Trie）

TreeNode.java

package Trie;
 
public class TreeNode {
 
    public TreeNode[] slot = new TreeNode[26];
 
    public char c;
 
    public boolean isWord;
 
    public int prefix;
 
    public String word;
 
    public String explain;
 
}

Trie.java

package Trie;
 
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
 
public class Trie {
 
    public static final TreeNode wordsTree = new TreeNode();
 
    public void insert(String words, String explain) {
        var root = wordsTree;
        var chars = words.toCharArray();
        for (var c : chars) {
            var idx = c - 'a';
            if (root.slot[idx] == null) {
                root.slot[idx] = new TreeNode();
            }
            root = root.slot[idx];
            root.c = c;
            root.prefix++;
        }
        root.explain = explain;
        root.isWord = true;
    }
 
    public List searchPrefix(String prefix){
        var root = wordsTree;
        var chars = prefix.toCharArray();
        var cache = new StringBuilder();
        for (var c : chars) {
            var idx = c - 'a';
            if (idx > root.slot.length || idx < 0 || root.slot[idx] == null) {
                return Collections.emptyList();
            }
            cache.append(c);
            root = root.slot[idx];
        }
        var list = new ArrayList();
        if (root.prefix != 0) {
            for (int i = 0; i < root.slot.length; i++) {
                if (root.slot[i] != null) {
                    var c = (char) (i + 'a');
                    collect(root.slot[i], String.valueOf(cache) + c, list, 15);
                    if (list.size() >= 15) {
                        return list;
                    }
                }
            }
        }
        return list;
    }
 
    protected void collect(TreeNode treeNode, String pre, List queue, int resultLimit) {
        if (treeNode.isWord) {
            treeNode.word = pre;
            queue.add(treeNode.word + "->" + treeNode.explain);
            if (queue.size() >= resultLimit) {
                return;
            }
        }
        for (int i = 0; i < treeNode.slot.length; i++) {
            var c = (char) ('a' + i);
            if (treeNode.slot[i] != null) {
                collect(treeNode.slot[i], pre + c, queue, resultLimit);
            }
        }
    }
 
}

测试代码

    public static void main(String[] args) {
        var trie = new Trie();
        trie.insert("bat", "大厂");
        trie.insert("batch", "批量");
        trie.insert("bitch", "彪子");
        trie.insert("battle", "战斗");
        System.out.println(trie);
        List trieNodes = trie.searchPrefix("ba");
        System.out.println(trieNodes);
    }