【算法】前缀树

概述

本质是一个多叉树，利用字符串的公共前缀来减少查询时间。查询、插入的时间复杂度为O(s) （s为操作的字符串长度）

上图为下面模板插入后生成的多叉树

模板

import java.util.HashMap;
import java.util.Map;
import java.util.Stack;

public class PrefixTree {
    static class Node {
        public Map<Character, Node> childrenMap; // 多路子树
        public boolean isEnd; // 是否是字符串结尾
        public Node() {
            childrenMap = new HashMap<>();
        }
    }
    // 整颗前缀树的根节点
    private final Node root;
    public PrefixTree() {
        root = new Node();
    }

    /**
     * 插入字符串，不存在返回true，存在返回false
     */
    public boolean insert(String s) {
        if (s == null) {
            return false;
        }
        Node tmp = root;
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(i);
            Node next = tmp.childrenMap.get(c);
            if (next == null) {
                next = new Node();
                tmp.childrenMap.put(c, next);
            }
            tmp = next;
        }
        if (tmp.isEnd) {
            return false;
        } else {
            tmp.isEnd = true;
            return true;
        }
    }

    /**
     * 判断是否完全匹配
     */
    public boolean contains(String s) {
        PrefixTree.Node res = searchPrefix(s);
        return res != null && res.isEnd;
    }

    /**
     * 是否前缀匹配
     */
    public boolean containsPrefix(String s) {
        return searchPrefix(s) != null;
    }

    /**
     * 删除字符串
     */
    public boolean delete(String s) {
        Node tmp = root;
        Stack<Node> stack = new Stack<>();
        for (int i = 0; i < s.length(); i++) {
            stack.push(tmp);
            char c = s.charAt(i);
            Node next = tmp.childrenMap.get(c);
            if (next == null) {
                return false;
            }
            tmp = next;
        }

        // 如果这个结点不是字符串结尾，说明树中不存在该字符串，返回false
        if (!tmp.isEnd) {
            return false;
        } else
        // 如果树中存在字符串，并且结点有子节点，说明不用从链表移除该节点， 将是否结尾置为false即可
        if (!tmp.childrenMap.isEmpty()) {
            tmp.isEnd = false;
            return true;
        } else // 树中存在该字符串，并且该节点是叶子结点
        {
            // 从下往上依次删除结点，如果某个结点是字符串结尾，或者某节点还有其它子节点，则不用删除这个结点，删除结束
            for (int i = s.length() - 1; i >= 0; i--) {
                Node next = stack.pop();
                next.childrenMap.remove(s.charAt(i));
                if (next.isEnd || !next.childrenMap.isEmpty()) {
                    return true;
                }
            }
        }
        // 走到这里说明删到了根部
        return true;
    }

    private Node searchPrefix(String s) {
        Node tmp = root;
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(i);
            Node next = tmp.childrenMap.get(c);
            if (next == null) {
                return null;
            }
            tmp = next;
        }
        return tmp;
    }

    public static void main(String[] args) {
        PrefixTree prefixTree = new PrefixTree();
        System.out.println("prefixTree.insert(\"abc\"):" + prefixTree.insert("abc"));
        System.out.println("prefixTree.insert(\"def\"):" + prefixTree.insert("def"));
        System.out.println("prefixTree.insert(\"abc\"):" + prefixTree.insert("abc"));
        System.out.println("prefixTree.insert(\"abd\"):" + prefixTree.insert("abd"));
        System.out.println("prefixTree.contains(\"abc\"):" + prefixTree.contains("abc"));
        System.out.println("prefixTree.contains(\"ab\"):" + prefixTree.contains("ab"));
        System.out.println("prefixTree.containsPrefix(\"ab\"):" + prefixTree.containsPrefix("ab"));
        System.out.println("prefixTree.containsPrefix(\"abcd\"):" + prefixTree.containsPrefix("abcd"));
        System.out.println("prefixTree.delete(\"abc\"):" + prefixTree.delete("abc"));
        System.out.println("prefixTree.contains(\"abc\"):" + prefixTree.contains("abc"));
        System.out.println("prefixTree.contains(\"ab\"):" + prefixTree.contains("ab"));
        System.out.println("prefixTree.containsPrefix(\"ab\"):" + prefixTree.containsPrefix("ab"));
        System.out.println("prefixTree.containsPrefix(\"abcd\"):" + prefixTree.containsPrefix("abcd"));
    }
}

输出结果

prefixTree.insert(“abc”):true
prefixTree.insert(“def”):true
prefixTree.insert(“abc”):false
prefixTree.insert(“abd”):true
prefixTree.contains(“abc”):true
prefixTree.contains(“ab”):false
prefixTree.containsPrefix(“ab”):true
prefixTree.containsPrefix(“abcd”):false
prefixTree.delete(“abc”):true
prefixTree.contains(“abc”):false
prefixTree.contains(“ab”):false
prefixTree.containsPrefix(“ab”):true
prefixTree.containsPrefix(“abcd”):false

练习题目：208. 实现 Trie (前缀树)