字典树[包含搜索和计算数量的功能]

package com.scene.field;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

/**
 * @Author luoshu
 * @Class TrieNode
 * @Description 字典树节点

 */
public class TrieNode {
    private Character data;
    private Map<Character, TrieNode> trieNodes;

    private Integer num = 0;
    private String res;


    public TrieNode(Character data) {
        this.data = data;
    }

    public TrieNode addOrGetTrieNode(Character c) {
        if (trieNodes == null) {
            trieNodes = new HashMap<>();
        }
        TrieNode trieNode = trieNodes.get(c);
        String s = Optional.ofNullable(res).orElse("");
        s += c;
        if (trieNode == null) {
            trieNode = new TrieNode(c);
            trieNode.res = s;
            trieNodes.put(c, trieNode);
        }
        return trieNode;
    }


    public String getRes() {
        return res;
    }

    public Character getData() {
        return data;
    }

    public Map<Character, TrieNode> getTrieNodes() {
        return trieNodes;
    }


    public Integer getNum() {
        return num;
    }

    public void addOne() {
        num++;
    }

    public void deleteOne() {
        if (num > 0) {
            num--;
        }
    }

    public void setNum(Integer num) {
        this.num = num;
    }
}

package com.scene.field;

import java.util.*;
import java.util.function.Function;

/**
 * @Author luoshu
 * @Class TrieTree
 * @Description 字典树

 */
public class TrieTree {
     private TrieNode root = new TrieNode('0');

    /**
     * 往树里添加词
     *
     * @param str
     */
    public void add(String str) {
        char[] chars = str.toCharArray();
        TrieNode t = root;
        for (char c : chars) {
            t = t.addOrGetTrieNode(c);
        }
        t.addOne();
    }

    public void delete(String str) {
        char[] chars = str.toCharArray();
        TrieNode t = root;
        TrieNode p = null;
        for (char c : chars) {
            p = t;
            t = t.addOrGetTrieNode(c);
        }
        //如果t没有子节点了 那可以把t直接移除了
        if (t.getTrieNodes() == null) {
            p.getTrieNodes().remove(t.getData());
        } else {
            t.deleteOne();
        }
    }

    /**
     * 通过前缀来搜索字典
     *
     * @param prefix 前缀
     * @param size   最大搜索数量  默认10
     * @return
     */
    public List<TrieNode> listTrieNode(String prefix, Integer size) {
        List<TrieNode> list = new LinkedList<>();
        if (size == null) {
            size = 10;
        }
        //进行BFS遍历
        Queue<TrieNode> queue = new LinkedList<>();
        queue.offer(root);
        //已经添加的结果数量
        int num = 0;
        //树的深度
        int deep = 0;
        while (!queue.isEmpty()) {
            //每次遍历深度+1
            int len = queue.size();
            for (int i = 0; i < len; i++) {
                TrieNode poll = queue.poll();
                //如果搜索数量已经够了 直接返回
                if (size <= num) {
                    break;
                }
                //只有该节点是一个单词 并且搜索深度已经大于前缀了的时候 才是正确的结果
                if (poll.getNum() > 0 && deep >= prefix.length()) {
                    list.add(poll);
                    num++;
                }
                //判断树的子节点是否存在
                if (poll.getTrieNodes() != null) {
                    for (Map.Entry<Character, TrieNode> entry : poll.getTrieNodes().entrySet()) {
                        //如果深度在小于搜索前缀的时候 要对当前节点和搜索前缀的该char进行对比 如果相同才入队
                        //深度大于前缀的时候直接入队 因为要搜索以词缀开头的子树
                        if (deep < prefix.length() && prefix.charAt(deep) != entry.getValue().getData()) {
                            continue;
                        }
                        queue.add(entry.getValue());
                    }
                }
            }
            deep++;
        }
        return list;
    }


    public List<String> list(String prefix, Integer size) {
        List<TrieNode> search = listTrieNode(prefix, size);
        List<String> list = new LinkedList<>();
        search.forEach(item -> {
            list.add(item.getRes());
        });
        return list;
    }


    public Integer count(String world) {

        List<TrieNode> search = listTrieNode(world, 1);
        return search.size() == 0 ? 0 : search.get(0).getNum();
    }

    public static void main(String[] args) {
        TrieTree trieTree=new TrieTree();
        trieTree.add("asdasda");
        trieTree.add("asdxasda");
        List<TrieNode> list=trieTree.search("asd",1);
        System.out.println(list);


    }
}