利用字典树实现搜索提示词

利用字典树构建搜索提示，主要一点是在项目中字典的数据加载。

    @GetMapping("/searchCues")
    @Operation(summary ="搜索提示词")
    public Result searchCues(String keys){
        try{
            if (keys.isEmpty()){
                return null;
            }
            return Result.success(trie.searchList(keys),"结果");
        }catch (Exception e){
            return Result.error(e.getMessage() + "搜索提示词");
        }
    }

import lombok.extern.slf4j.Slf4j;
import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.UnsupportedEncodingException;
/**
 * 加载词典用的类
 * @author sol
 * @date 2023/9/4 10:11
 * @Version 1.0
 */
@Slf4j
public class DicReader {
    public static BufferedReader getReader(String name) {
        // maven工程修改词典加载方式
        InputStream in = DicReader.class.getResourceAsStream(name);
        try {
            return new BufferedReader(new InputStreamReader(in, "UTF-8"));
        } catch (UnsupportedEncodingException e) {
            log.warn("不支持的编码", e);
        }
        return null;
    }
    public static InputStream getInputStream(String name) {
        // maven工程修改词典加载方式
        InputStream in = DicReader.class.getResourceAsStream(name);
        return in;
    }
}

初始化字典

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.UnsupportedEncodingException;

/**
 * @author sol
 * @date 2023/9/4 10:22
 * @Version 1.0
 */
@Component
@Slf4j
public class DicInit {
    Trie chTrie = new Trie();
    public  DicInit(){

　　　　　// 字典加载路径 项目resources根目录下 library/default.dic 其中default是你自己的提示词词典

try (BufferedReader reader = DicReader.getReader("/library/default.dic")) {
            String temp = null;
            while ((temp = reader.readLine()) != null) {
                if (temp.isEmpty()) {
                    continue;
                }
                chTrie.insert(temp);
            }
        } catch (UnsupportedEncodingException e) {
            log.warn("不支持的编码", e);
        } catch (IOException e) {
            log.warn("IO异常", e);
        }
    }
    public Trie getTrie(){
        return this.chTrie;
    }
}

字典树结构

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author sol
 * @date 2023/9/2 15:44
 * @Version 1.0
 */
public class Trie implements Serializable {

    public static TrieNode root ;
    static class TrieNode {
        // 标识是否为终节点
        Boolean isEnd;
        // 多少个路线通过该节点，（有多少子节点 + 1）？
        Integer num;
        // 保存所有儿子节点
        Map<String,TrieNode> sonMap;
        // 标记是否存在儿子节点
        Boolean hasSon;

        TrieNode() {
            isEnd = false;
            hasSon = false;
            num = 1;
            sonMap = new HashMap<String, TrieNode>();
        }
    }

    public Trie() {
        if (root == null) {
            root = new TrieNode();
        }
    }
    /**
     * 插入字典树
     * @param sentence
     */
    public void insert(String sentence) {
        TrieNode node = root;
        for (int i=0; i<sentence.length(); i++) {
            String word = String.valueOf(sentence.charAt(i));
            // 如果树中不存在该字
            if (!node.sonMap.containsKey(word)) {
                node.sonMap.put(word,new TrieNode());
                node.hasSon = true;
            } else {
                node.num ++;
            }
            node = node.sonMap.get(word);
        }
        node.isEnd = true;
    }

    /**
     * 在树中查找
     * @param sentence
     * @return
     */
    public boolean search(String sentence) {
        TrieNode node = root;
        for (int i=0; i<sentence.length(); i++) {
            String word = String.valueOf(sentence.charAt(i));
            if (node.sonMap.containsKey(word)) {
                node = node.sonMap.get(word);
            } else {
                return false;
            }
        }
        return node.isEnd;
    }

    /**
     * 获取所有以入参关键字开头的词汇
     * @param sentence
     * @return
     */
    public List<String> searchList(String sentence) {
        List<String> resultList = new ArrayList<>();
        TrieNode node = root;
        StringBuilder prefix = new StringBuilder();
        for (int i=0; i<sentence.length(); i++) {
            String word = String.valueOf(sentence.charAt(i));
            // 如果树中不包含关键字，直接返回
            if (!node.sonMap.containsKey(word)) {
                return resultList;
            } else {
                node = node.sonMap.get(word);
            }
            prefix.append(word);
            // 判断如果是一句话中的最后一个字符
            if (i == sentence.length() -1) {
                if (search(sentence)){
                    resultList.add(prefix + "");
                }
                recursion(node, resultList, prefix.toString());
            }
        }
        return resultList;
    }

    /**
     * 递归获取字典树的树枝。
     * @param node
     * @return
     */
    private String recursion(TrieNode node, List<String> strings,  String sb) {
        if (node.hasSon) {
            for (String key : node.sonMap.keySet()) {
                if (node.sonMap.get(key).isEnd) {
                    strings.add(sb+key);
                }
            }
            for (String key : node.sonMap.keySet()) {
                recursion(node.sonMap.get(key),strings,sb+key);
            }
        }
        return sb.toString();
    }
}