【前缀树算法】java超详细！一文带你手撕前缀树

原创已于 2024-06-28 14:04:04 修改 · 714 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#算法 #java #开发语言

于 2024-06-18 10:42:38 首次发布

算法笔记专栏收录该内容

2 篇文章

订阅专栏

参考 b站左程云讲解045视频
配合视频使用更佳！

本文通过两种方式实现前缀树，并给出了牛客测试链接

最后通过这两种方法完成例题【接头密匙】

类描述方式实现前缀树

路的可能性范围较小，用固定数组实现路 path
路的可能性范围较大，用哈希表实现路 path

固定数组

假设这里字符串由 26 个英文字母组成，可以使用固定数组-类描述方式实现

public class TrieTree {
    class TrieNode{
        private int pass;
        private int end;
        private TrieNode[] nexts;
        TrieNode(){
            pass = 0;
            end = 0;
            nexts = new TrieNode[26];
        }
    }

    private TrieNode root;

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

    public void insert(String word){
        TrieNode node = root;
        node.pass++;
        for(int i = 0, path; i < word.length(); i++){
            path = word.charAt(i) - 'a';
            if(node.nexts[path] == null){
                node.nexts[path] = new TrieNode();
            }
            node = node.nexts[path];
            node.pass++;
        }
        node.end++;
    }

    /**
     * 查询前缀树中word单词出现了几次
     * @param word
     * @return
     */
    public int countWords(String word){
        TrieNode node = root;
        for(int i = 0, path; i < word.length(); i++){
            path = word.charAt(i) - 'a';
            //如果没有 说明没出现过
            if(node.nexts[path] == null){
                return 0;
            }
            node = node.nexts[path];
        }
        return node.end;
    }

    /**
     * 查询以pre为前缀的单词有多少个
     * @param pre
     * @return
     */
    public int countPrefix(String pre){
        TrieNode node = root;
        for(int i = 0, path; i < pre.length(); i++){
            path = pre.charAt(i) - 'a';
            if(node.nexts[path] == null){
                return 0;
            }
            node = node.nexts[path];
        }
        return node.pass;
    }

    /**
     * 删除单词word
     * 如果没出现过，就什么也不做
     * @param word
     */
    public void erase(String word){
        if(countWords(word) > 0){
            TrieNode node = root;
            node.pass--;
            for(int i = 0, path; i < word.length(); i++){
                path = word.charAt(i) - 'a';
                if(--node.nexts[path].pass == 0){
                    //pass--后为0了，就可以把后续的删掉了
                    node.nexts[path] = null;
                    return;
                }
                node = node.nexts[path];
            }
            node.end--;
        }
    }
}

哈希表

假设这里字符串由 int 组成，使用哈希表-类描述方式实现

public class TrieTree2 {
    class TrieNode{
        public int pass;
        public int end;
        HashMap<Integer, TrieNode> nexts;

        public TrieNode(){
            pass = 0;
            end = 0;
            nexts = new HashMap<>();
        }
    }

    private TrieNode root;

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

    public void insert(String word){
        TrieNode node = root;
        node.pass++;
        for (int i = 0, path; i < word.length(); i++) {
            path = word.charAt(i);
            if(!node.nexts.containsKey(path)){
                node.nexts.put(path, new TrieNode());
            }
            node = node.nexts.get(path);
            node.pass++;

        }
        node.end++;
    }

    public int countWords(String word){
        TrieNode node = root;
        for (int i = 0, path; i < word.length(); i++) {
            path = word.charAt(i);
            if (!node.nexts.containsKey(path)) {
                return 0;
            }
            node = node.nexts.get(path);
        }
        return node.end;
    }

    public int countPrefix(String pre){
        TrieNode node = root;
        for (int i = 0, path; i < pre.length(); i++) {
            path = pre.charAt(i);
            if (!node.nexts.containsKey(path)) {
                return 0;
            }
            node = node.nexts.get(path);
        }
        return node.pass;
    }

    public void erase(String word){
        if (countWords(word) > 0) {
            TrieNode node = root;
            TrieNode next;
            node.pass--;
            for (int i = 0, path; i < word.length(); i++) {
                path = word.charAt(i);
                next = node.nexts.get(path);
                if (--next.pass == 0) {
                    node.nexts.remove(path);
                    return;
                }
                node = next;
            }
            node.end--;
        }
    }
}

静态数组方式实现前缀树

https://www.nowcoder.com/practice/7f8a8553ddbf4eaab749ec988726702b

使用一个 cnt 变量表示目前树中记录到第几个结点，从 1 开始

用二维静态数组表示树，tree[][] 第一位表示第几个结点，第二位表示该结点都连了哪些结点

tree[cur][path] 为 0 时表示没有

每次新建结点时 ++cnt

删除节点时，如果遇到某一处 pass 为 0，则直接将该处的 tree 置为 0 即可，后面的丢弃掉，这样就算再新建树时，也是++cnt，不会用到丢弃掉的数

    public static int MAXN = 150001;
    public static int[][] tree = new int[MAXN][26];
    public static int[] pass = new int[MAXN];
    public static int[] end = new int[MAXN];

    public static int cnt;

    public static void build(){
        cnt = 1;
    }

    public static void insert(String word){
        int cur = 1;
        pass[cur]++;
        for (int i = 0, path; i < word.length(); i++) {
            path = word.charAt(i) - 'a';
            if(tree[cur][path] == 0){
                tree[cur][path] = ++cnt;
            }
            cur = tree[cur][path];
            pass[cur]++;
        }
        end[cur]++;
    }

    public static int search(String word){
        int cur = 1;
        for(int i = 0, path; i < word.length(); i++){
            path = word.charAt(i) - 'a';
            if(tree[cur][path] == 0){
                return 0;
            }
            cur = tree[cur][path];
        }
        return end[cur];
    }

    public static int prefixNumber(String word){
        int cur = 1;
        for(int i = 0, path; i < word.length(); i++){
            path = word.charAt(i) - 'a';
            if(tree[cur][path] == 0){
                return 0;
            }
            cur = tree[cur][path];
        }
        return pass[cur];
    }

    public static void delete(String word){
        if(search(word) <= 0){
            return;
        }
        int cur = 1;
        for (int i = 0, path; i < word.length(); i++) {
            path = word.charAt(i) - 'a';
            if(--pass[tree[cur][path]] == 0){
                tree[cur][path] = 0;
                return;
            }
            cur = tree[cur][path];
        }
        end[cur]--;
    }

    public static void clear(){
        for (int i = 0; i < cnt; i++) {
            Arrays.fill(tree[i], 0);
            pass[i] = 0;
            end[i] = 0;
        }
    }

public static int m, op;
    public static String[] splits;
    public static String line;

    public static void main(String[] args) throws IOException {
        BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
        PrintWriter out = new PrintWriter(new OutputStreamWriter(System.out));
        while((line = in.readLine()) != null){
            build();
            m = Integer.valueOf(line);
            for (int i = 1; i <= m; i++) {
                splits = in.readLine().split(" ");
                op = Integer.valueOf(splits[0]);
                if(op == 1){
                    insert(splits[1]);
                }else if (op == 2){
                    delete(splits[1]);
                }else if (op == 3){
                    out.println(search(splits[1]) > 0 ? "YES" : "NO");
                } else if (op == 4) {
                    out.println(prefixNumber(splits[1]));
                }
            }
            clear();
        }
        out.flush();
        in.close();
        out.close();
    }

例题 - 接头密匙

https://www.nowcoder.com/practice/c552d3b4dfda49ccb883a6371d9a6932

静态数组

较难理解，建议参考b站左程云讲解045视频
使用哈希表方法理解较简单

[2, 14, 6, 8] 则计算完后为 [12, -8, 2]，这里将其拼接为字符串"12#-8#2#"
#代表一个数字的结尾
将该字符串插入前缀树：
在这里插入图片描述

public class TrieTree_Code01 {
    public static int MAXN = 200000;
    public static int[][] tree = new int[MAXN][12];
    public static int[] pass = new int[MAXN];
    public static int[] end = new int[MAXN];

    public static int cnt;

    public static void build(){
        cnt = 1;
    }

    public static void clear(){
        for (int i = 1; i <= cnt; i++) {
            pass[i] = 0;
            end[i] = 0;
            Arrays.fill(tree[i], 0);
        }
    }

    public static int path(char c){
        if(c == '#'){
            return 10;
        } else if (c == '-') {
            return 11;
        }else {
            return c - '0';
        }
    }

    public static void insert(String word){
        int cur = 1;
        pass[cur]++;
        for (int i = 0, path; i < word.length(); i++) {
            path = path(word.charAt(i));
            if(tree[cur][path] == 0){
                tree[cur][path] = ++cnt;
            }
            cur = tree[cur][path];
            pass[cur]++;
        }
        end[cur]++;
    }

    public static int searchCount(String word){
        int cur = 1;
        for (int i = 0, path; i < word.length(); i++) {
            path = path(word.charAt(i));
            if(tree[cur][path] == 0){
                return 0;
            }
            cur = tree[cur][path];
        }
        return pass[cur];
    }


    public int[] countConsistentKeys (int[][] b, int[][] a){
        build();
        StringBuilder sb = new StringBuilder();
        //将a添加到前缀树中
        for(int[] nums : a){
            sb.setLength(0);
            for(int i = 1; i < nums.length; i++){
                sb.append(String.valueOf(nums[i] - nums[i - 1]) + '#');
            }
            insert(sb.toString());
        }
        int[] ans = new int[b.length];
        for(int i = 0; i < b.length; i++){
            sb.setLength(0);
            int[] nums = b[i];
            for(int j = 1; j < nums.length; j++){
                sb.append(String.valueOf(nums[j] - nums[j - 1]) + '#');
            }
            ans[i] = searchCount(sb.toString());
        }
        clear();
        return ans;
    }
}

哈希表

直接用哈希表的方式，不拼接了
[2, 14, 6, 8] 则计算完后为 [12, -8, 2]，将该数组插入前缀树
在这里插入图片描述

public class TrieTree_Code01_2 {
    static class TrieNode{
        public int pass;
        public int end;
        public HashMap<Integer, TrieNode> nexts;
        TrieNode(){
            pass = 0;
            end = 0;
            nexts = new HashMap<>();
        }
    }
    private static TrieNode root = new TrieNode();

    public static void insert(int[] word){
        TrieNode node = root;
        node.pass++;
        for(int i = 0; i < word.length; i++){
            int path = word[i];
            if(!node.nexts.containsKey(path)){
                node.nexts.put(path, new TrieNode());
            }
            node = node.nexts.get(path);
            node.pass++;
        }
        node.end++;
    }

    public static int searchPrefix(int[] pre) {
        TrieNode node = root;
        for (int i = 0; i < pre.length; i++) {
            int path = pre[i];
            if(!node.nexts.containsKey(path)){
                return 0;
            }
            node = node.nexts.get(path);
        }
        return node.pass;
    }

    public static int[] countConsistentKeys (int[][] b, int[][] a) {
        for(int[] nums : a){
            int[] word = new int[nums.length - 1];
            for(int i = 1; i < nums.length; i++){
                word[i - 1] = nums[i] - nums[i - 1];
            }
            insert(word);
        }

        int[] ans = new int[b.length];

        for(int i = 0; i < b.length; i++){
            int[] nums = b[i];
            int[] word = new int[nums.length - 1];
            for(int j = 1; j < nums.length; j++){
                word[j - 1] = nums[j] - nums[j - 1];
            }
            ans[i] = searchPrefix(word);
        }
        return ans;
    }

    public static void main(String[] args) {
        int[][] b = {{1, 2, 3, 4, 5}, {2, 4, 6, 8}, {1, 4, 7, 10}};
        int[][] a= {{3, 4, 5, 6, 7, 8}, {2, 4, 6, 8}, {1, 3, 5, 7, 9}};
        int[] ans = countConsistentKeys(b, a);
        for (int i = 0; i < ans.length; i++) {
            System.out.println(ans[i]);
        }
    }
}