【使用Trie树实现自动搜索完成（联想查询）】——升级搜索体验：Trie树助力智能搜索联想，助您轻松实现自动完成

纪大侠

已于 2024-01-15 14:34:27 修改

阅读量1k

点赞数 23

文章标签： java 数据库搜索引擎

于 2024-01-15 14:30:22 首次发布

本文链接：https://blog.csdn.net/weixin_42887222/article/details/135599788

版权

前言

自动搜索完成（联想查询），就是我们使用百度或者谷歌搜索引擎时，会出现候选项目，让我们可以快速选择搜索。例如：
在这里插入图片描述

现在基本上很多查询都有这类似的功能，已经是标配了。不过在十几年前还是很少有这种比较友好的查询方式，这边我们用一种叫Trie树数据结构实现联想查询。

自动搜索完成（Autocomplete）是一种用户界面设计技术，它在用户输入时自动提供与输入相关的建议或完成的搜索词汇。这有助于用户更快地找到他们正在搜索的内容，提高搜索的效率和用户体验。

Treie树

Trie树（前缀树、字典树）
Trie树常用于字符串的插入、删除、查找等操作。它特别适用于搜索引擎、拼写检查、自动完成等应用，因为它可以高效地存储和检索大量字符串。

Trie（发音为 “try”）的名称来源于英语单词 “retrieval”（检索）

代码实现

public class Trie {
    private final TrieNode root;

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

    public void insert(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (node.getChild(ch) == null) {
                node.setChild(ch, new TrieNode());
            }
            node = node.getChild(ch);
        }
        node.setEndOfWord(true);
    }

    public boolean search(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (node.getChild(ch) == null) {
                return false;
            }
            node = node.getChild(ch);
        }
        return node.isEndOfWord();
    }

    public boolean startsWith(String prefix) {
        TrieNode node = root;
        for (int i = 0; i < prefix.length(); i++) {
            char ch = prefix.charAt(i);
            if (node.getChild(ch) == null) {
                return false;
            }
            node = node.getChild(ch);
        }
        return true;
    }

    public static void main(String[] args) {
        Trie trie = new Trie();

        trie.insert("apple");
        System.out.println(trie.search("apple")); // 输出 true
        System.out.println(trie.search("app"));   // 输出 false
        System.out.println(trie.startsWith("app")); // 输出 true

        trie.insert("app");
        System.out.println(trie.search("app"));   // 输出 true
    }
}

class TrieNode {
    private final TrieNode[] children;
    private boolean isEndOfWord;

    public TrieNode() {
        children = new TrieNode[26]; // 假设只包含小写字母
        isEndOfWord = false;
    }

    public TrieNode getChild(char ch) {
        return children[ch - 'a'];
    }

    public void setChild(char ch, TrieNode node) {
        children[ch - 'a'] = node;
    }

    public boolean isEndOfWord() {
        return isEndOfWord;
    }

    public void setEndOfWord(boolean endOfWord) {
        isEndOfWord = endOfWord;
    }


}

前缀树问题

从上面看到前缀树，每个节点的子节点只有26个，以上只支持输入的字符是全是小写字母。这个在我们现实世界中字符光中文就有几万个，每添加一个词语，那么就要增加一个几万个数组长度的。这边会造成很多的性能浪费。那么就换个方式用hash表进行映射。

Trie树 vs 倒排索引

Trie 树：

设计思想： Trie 树是一种树形结构，每个节点表示一个字符，从根节点到某个节点的路径构成一个字符串。它的设计主要用于高效地支持前缀匹配。

适用场景： Trie 树适合处理需要高效前缀匹配的场景，例如自动补全、单词搜索等。通过 Trie 树，可以快速找到具有相同前缀的所有字符串。

优点：高效的前缀匹配，支持动态插入和删除，无歧义匹配。

缺点：占用空间较大，不适用于大规模文本数据。

倒排索引：

设计思想：倒排索引是一种以单词为关键字构建的数据结构，它将文档集合中每个单词映射到包含该单词的文档列表。主要用于支持文本搜索和检索。

适用场景：倒排索引适合处理大规模文本数据，用于构建搜索引擎或实现复杂的文本检索需求。通过倒排索引，可以快速找到包含特定单词的文档。

优点：高效的文本搜索，支持模糊搜索、通配符查询等。适用于大规模文本数据。

缺点：插入和删除效率较低，不适用于实时更新的场景

其中目前比较火的搜索引擎Elasticsearch 使用倒排索引（Inverted Index）作为其主要数据结构。

联想查询实现

要实现联想查询功能，用户输入一个词，后台返回频率最高的五个相关字符。为了提高查询速度，我们希望在 Trie 树的每个节点缓存五个频率最高的字符。当查询到前缀时，可以直接返回缓存的结果。我们将词语的数据源来自ansj分词-分仓库，数据量达到十万以上，足够作为演示。

这个优化方案将通过 Trie 树实现联想查询，并在 Trie 树的每个节点上保存频率最高的字符信息，以提高查询性能。同时，数据源的选取来自一个包含大量实际词汇的仓库，使得演示更具实际效果。

代码实现

package demo.algorithm.algorithm.sdi;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.*;

/**
 * 使用前缀树，检索单词。每个节点返回缓存的5个高频单词
 * 实际自动完成，会对英文全部转成小写，这样不区分大小写匹配。
 *
 * @author jisl on 2024/1/12 16:01
 **/
class WordInfo {
    private String word;
    private int frequency;

    public WordInfo(String word, int frequency) {
        this.word = word;
        this.frequency = frequency;
    }

    public String getWord() {
        return word;
    }

    public void setWord(String word) {
        this.word = word;
    }

    public int getFrequency() {
        return frequency;
    }
}

class TrieNode {
    private Map<Character, TrieNode> children;
    private boolean isEndOfWord;
    // 保存单词和频率信息的列表
    private List<WordInfo> wordList;

    public TrieNode() {
        children = new HashMap<>();
        isEndOfWord = false;
        wordList = new ArrayList<>();
    }

    public Map<Character, TrieNode> getChildren() {
        return children;
    }

    //    child单个儿子，children多个儿子
    public TrieNode getChild(char ch) {
        return children.get(ch);
    }

    public boolean isEndOfWord() {
        return isEndOfWord;
    }

    public void setEndOfWord(boolean endOfWord) {
        isEndOfWord = endOfWord;
    }

    public List<WordInfo> getWordList() {
        return wordList;
    }
}

class Trie {
    private TrieNode root;

    private final int k = 10;

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

    public void insert(String word, int frequency) {
        TrieNode node = root;
        for (char ch : word.toLowerCase().toCharArray()) {
            node.getChildren().putIfAbsent(ch, new TrieNode());
            node = node.getChild(ch);
            insertIntoWordList(node, word, frequency);
        }
        node.setEndOfWord(true);
    }

    private void insertIntoWordList(TrieNode node, String word, int frequency) {
        List<WordInfo> wordList = node.getWordList();
        wordList.add(new WordInfo(word, frequency));
        Collections.sort(wordList, (w1, w2) -> Integer.compare(w2.getFrequency(), w1.getFrequency()));
        // 保留前k个单词
        if (wordList.size() > k) {
            wordList.remove(k);
        }
    }

    public List<WordInfo> searchTopWords(String prefix) {
        TrieNode node = root;
        for (char ch : prefix.toLowerCase().toCharArray()) {
            node = node.getChild(ch);
            if (Objects.isNull(node)) {
                return Collections.emptyList();  // 没有以该前缀开始的单词
            }
        }
        return node.getWordList();
    }
}

public class TrieWithFrequencyExample {

    private static List<WordInfo> readWordInfoFromFile(String filePath) {
        List<WordInfo> wordInfoList = new ArrayList<>();

        try (BufferedReader br = new BufferedReader(new FileReader(filePath))) {
            String line;
            while ((line = br.readLine()) != null) {
                String[] parts = line.split("\t");
                if (parts.length == 3) {
                    String word = parts[0].trim();
//                    String pos = parts[1].trim();
                    int frequency = Integer.parseInt(parts[2].trim());

                    WordInfo wordInfo = new WordInfo(word, frequency);
                    wordInfoList.add(wordInfo);
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return wordInfoList;
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        String path = System.getProperty("user.dir") + "\\files\\ansj_seg\\default.dic";
        final List<WordInfo> wordInfoList = TrieWithFrequencyExample.readWordInfoFromFile(path);
        for (WordInfo wordInfo : wordInfoList) {
            trie.insert(wordInfo.getWord(), wordInfo.getFrequency());
        }
        System.out.println("单词数量:" + wordInfoList.size());
        trie.insert("Apple", 19);
        trie.insert("apricot", 8);
        trie.insert("apex", 15);
        trie.insert("bat", 5);
        trie.insert("batman", 12);
        trie.insert("batwoman", 7);
        trie.insert("banana", 20);
        trie.insert("app", 18);
        trie.insert("apartment", 9);
        trie.insert("appliance", 11);
        trie.insert("李白", 8);
        trie.insert("李太白", 8);
        trie.insert("李白诗", 7);
        trie.insert("李白 侠客行", 6);
        trie.insert("李白代表作", 5);
        trie.insert("李白 杜甫", 5);
        // 查询以 "ap" 开头的前5个频率最高的单词
        String prefix = "长安";
        List<WordInfo> result = trie.searchTopWords(prefix);

        System.out.printf("Top Words with Prefix '%s':%n", prefix);
        for (WordInfo wordFreq : result) {
            System.out.println(wordFreq.getWord() + " - Frequency: " + wordFreq.getFrequency());
        }
    }
}

运行效果

单词数量:386260
Top Words with Prefix '长安':
长安 - Frequency: 5355
长安公司 - Frequency: 3429
长安大戏院 - Frequency: 3386
长安街 - Frequency: 1482
长安路 - Frequency: 1344

总结

以上实现思路是来自《System Design Interview》这本书里“CHAPTER 13: DESIGN A SEARCH AUTOCOMPLETE SYSTEM”实现一个自动搜索完成，为了加深印象那么就用代码实现下。刚开始想着使用倒排索引，因为中文的缘故。但是前缀搜索，倒排索引实现不是很好。