联想搜索实现方案（java基于字典树算法的实现方式）

1、背景

最近我的自动化测试平台（PostGirl）上有一个小需求：
用户在知识库的搜索框输入关键字，下方自动显示出以该关键字开头的词汇。实现效果类似百度的联想搜索（见下图）。

2、方案一

开始我的实现思路是使用redis的zset来实现。通过zadd添加元素。搜索的时候使用zrank获取到关键字的位置，然后通过zrange 得到所有以关键字开头的词汇，最后进行展示。
核心代码如下：

 // 1、将关键字存储到redis中，用于后续输入联想功能
 jedis.zadd("search",0,word.getName());

// 2、查询出关键字处于有序集合的位置
 Long wordIndex = resource.zrank("search", keyWord);
// 3、获取指定起始位置到末尾的所有值 
 Set<String> results = resource.zrange("search", wordIndex, -1L);
        for (String result: results) {
            if (result.startsWith(keyWord)) {
                reustlts.add(result); // 这个里面存储了所有以指定关键字开头的词汇。
            }
        }

但是，使用redis实现有个缺点，在zset中获取以某个字开头的位置很好确定，但是结尾不好确定，容易把不相关的内容搜索出来，感觉不怎么好处理。

方案二

在内存中维护一颗字典树，每次插入关键字后将字典树序列化为json字符保存到数据库。同时更新字典树对象。当重启的时候，将数据库中的json字符串查询出来然后实例化为对象（这一步在启动时加载效果更好）。
代码实现如下：

package com.cz.commons.search;

import com.cz.commons.holder.Holder;
import com.google.gson.Gson;

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

/**
 * 字典树，用于搜索联想
 * @program: PostGirl-panent
 * @description: Trie
 * @author: Cheng Zhi
 * @create: 2021-07-10 12:36
 **/
public class Trie implements Serializable {

    public static Trie trie;
    public static TrieNode root ;

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

    public static Trie getTrie() {
        if (trie == null) {
            trie = new Trie();
        }
        return trie;
    }

    static class TrieNode {

        // 标识是否为终节点
        Boolean isEnd;
        // 多少个路线通过该节点
        Integer num;
        // 保存所有儿子节点
        Map<String,TrieNode> sonMap;
        // 标记是否存在儿子节点
        Boolean hasSon;

        TrieNode() {
            isEnd = false;
            hasSon = false;
            num = 1;
            sonMap = new HashMap<String, 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();
        String[] strings = new String[20];
        int index = 0;
        Holder<Integer> holder = new Holder<>(index);
        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) {
                final String finalSb = prefix.toString();
                recursion(node, resultList, holder, prefix.toString());
            }
        }
        return resultList;
    }

    /**
     * 递归获取字典树的树枝。
     * @param node
     * @return
     */
    private String recursion(TrieNode node, List<String> strings, Holder<Integer> holder, String sb) {

        String prefix = sb;

        if (node.hasSon) {
            for (String key : node.sonMap.keySet()) {
                if (node.sonMap.get(key).isEnd) {
                    //strings[holder.getValue()] = sb+key;
                    strings.add(sb+key);
                }
                holder.setValue(holder.getValue() + 1);
            }
            for (String key : node.sonMap.keySet()) {
                holder.setValue(holder.getValue() + 1);
                recursion(node.sonMap.get(key),strings,holder, sb+key);
            }
        }
        return sb.toString();
    }

    /**
     * 为Trie内部类TrieNode初始化值
     * @param jsonObject
     * @return
     */
    public Trie setTrieNode(String jsonObject) {
        TrieNode trieNode = mapToObject(jsonToMap(jsonObject));
        Trie trie = getTrie();
        trie.root = trieNode;
        return trie;
    }

    /**
     * 将json转化为Map
     * @param jsonStr
     * @return
     */
    public Map<?, ?> jsonToMap(String jsonStr) {
        Map<?, ?> ObjectMap = null;
        Gson gson = new Gson();
        java.lang.reflect.Type type = new com.google.gson.reflect.TypeToken<Map<?,?>>() {}.getType();
        ObjectMap = gson.fromJson(jsonStr, type);
        return ObjectMap;
    }

    /**
     * 将map转化为java对象
     * @param map
     * @return
     */
    public TrieNode mapToObject(Map<?, ?> map) {

        TrieNode trieNode = new TrieNode();
        Double doubleNum = (Double) map.get("num");
        trieNode.num = doubleNum.intValue();;
        trieNode.isEnd = (Boolean) map.get("isEnd");
        trieNode.hasSon = (Boolean) map.get("hasSon");
        Map<?,?> sonMap = (Map<?, ?>) map.get("sonMap");
        Map<String, TrieNode> newMap = new HashMap<>();
        for (Object key : sonMap.keySet()) {
            newMap.put((String) key, mapToObject((Map<?, ?>) sonMap.get(key)));
        }
        trieNode.sonMap = newMap;
        return trieNode;
    }

    public static void main(String[] args) {
        Trie chTrie = new Trie();
        Gson gson = new Gson();
        String str = "这是一坨,这是,这是一箱,测试,这是个好孩子,这测试,这是行,这是好,这是啥,这是朱,这是猪啊啊啊啊啊啊啊啊";
        String[] splits = str.split(",");
        for (String split : splits) {
            chTrie.insert(split);
        }

        String strs = "这是";
        System.out.println(chTrie.searchList(strs));
    }
}

效果

思考

这样做，数据量小貌似没有问题，但是，如果数据量相当大的话，那么这颗树将会很大，使用数据库操作必定会很慢。还得再想想怎么搞。。。

补充

好多朋友询问Holder是什么？holder是一个用来做数据传递的载体，可以理解为是一个容器。
下面是源码：

package pers.cz.commons.holder;

/**
 * @program: netDisc  使用引用传递的方式实现不可变类的传递：例如int类型的传递
 * @description: Holder
 * @author: Cheng Zhi
 * @create: 2020-12-05 13:40
 **/
public class Holder<T> implements IHolder {

    private T value;
    private IHolder pipeHolder;

    public T getValue() {
        return value;
    }

    public void setValue(T value) {
        this.value = value;
    }

    public Holder(T value) {
        this.value = value;
    }

    public Holder() {

    }
    @Override
    public String toString() {
        return "Holder{" +
                "value=" + value +
                ", pipeHolder=" + pipeHolder +
                '}';
    }

    @Override
    public Object get() {
        return null;
    }

    @Override
    public void set(Object obj) {

    }

    /**
     * 测试方法
     * @param args
     */
    public static void main(String[] args) {
        // doing();
        String name  = "cz";
        Holder holder = new Holder(name);
        change(holder);
        System.out.println(holder.getValue());

        Integer a = 5;
        Holder holder1 = new Holder(a);
        changeInt(holder1);
        System.out.println(holder1.getValue());

    }

    private static void changeInt(Holder holder) {
        Integer a = (Integer) holder.getValue();
        holder.setValue( a+1);
    }
    private static void change(Holder holder) {
        holder.setValue("cz666");

    }
}

package pers.cz.commons.holder;

import java.io.Serializable;

public interface IHolder extends Serializable {

    public abstract Object get();

    public abstract void set(Object obj);

}