算法篇五 Trie树与AC自动机

public class LeetCode {
    public static void main(String[] args) {
        // 游戏开始, 寻找J老师
        Trie trie = new Trie();
        trie.insert("johonsona".toCharArray());
        trie.insert("johonsonb".toCharArray());
        trie.insert("johonson".toCharArray());
        trie.insert("yinan".toCharArray());
        trie.buildFailPointer();
        trie.match("adadajohonsonadada".toCharArray());
    }
}

class Trie {
    // root节点
    private TrieNode root = new TrieNode('/');

    // 纯Trie树: 往树中插入字符串
    public void insert(char[] text) {
        TrieNode cur = root;
        for (int i = 0; i < text.length; i++) {
            // a的ascii值为97, 这边只考虑26个小写字母的情况, 所以任何字母的下标都要先减去a
            // 比如 text[i] 为 b, 则 b - a = 98 - 97 = 1; 所以b的index就是1
            int index = text[i] - 'a';
            if (cur.children[index] == null) {
                // 初始化子树
                TrieNode newNode = new TrieNode(text[i]);
                cur.children[index] = newNode;
            }
            cur = cur.children[index];
        }
        // 在字符串每个字符全都插入之后, 把当前节点标记为一个结束节点
        cur.isEndingChar = true;
        cur.length = text.length;
    }

    // 纯Trie树: 从树中查找一个字符串
    public boolean find(char[] pattern) {
        TrieNode cur = root;
        for (int i = 0; i < pattern.length; i++) {
            int index = pattern[i] - 'a';
            if (cur.children[index] == null) {
                return false;
            }
            cur = cur.children[index];
        }
        return cur.isEndingChar;
    }

    // AC自动机: DFS构建失败指针
    public void buildFailPointer() {
        Queue<TrieNode> queue = new LinkedList<>();
        root.fail = null;
        queue.offer(root);
        while (!queue.isEmpty()) {
            TrieNode cur = queue.poll();
            for (int i = 0; i < 26; i++) {
                // 循环取出子树
                TrieNode curChild = cur.children[i];
                if (curChild == null) {
                    continue;
                }
                // 如果cur为root, 则子节点的失败指针全部指向root, 因为子节点就是第一层
                if (cur == root) {
                    curChild.fail = root;
                } else {
                    TrieNode curFail = cur.fail;
                    while (curFail != null) {
                        // 把这个子节点的失败指针指向上一层失败指针的子节点中 data相等的子节点
                        // 比如当前节点父节点为 a, 当前节点为 b, 此时 curChild.data 为 c
                        // 则根据 最长匹配后缀子串, 我们希望寻找 bc这个字符串, cur.fail肯定也是b(或者root)
                        // 所以 curChild.fail = cur.fail.children[c所在的index]
                        TrieNode curFailChild = curFail.children[curChild.data - 'a'];
                        if (curFailChild != null) {
                            curChild.fail = curFailChild;
                            break;
                        }
                        // 否则继续往上一层找, 比如 abc, 没有找到bc, 则直接找c
                        curFail = curFail.fail;
                    }
                    // 如果一直搜索到失败指针为null也没找到匹配的, 则指向root
                    // 这边判断用 if (curChild.fail == null) 也是一样的
                    if (curFail == null) {
                        curChild.fail = root;
                    }
                }
                queue.add(curChild);
            }
        }
    }

    // AC自动机: 查找所有存在trie树中的子串
    public void match(char[] text) {
        TrieNode cur = root;
        for (int i = 0; i < text.length; i++) {
            int index = text[i] - 'a';
            // 如果当前路径下匹配不到, 则直接换个路径, 借助失败指针来跳转
            // 比如当前路径是abc,text[i]为c, 如果找不到, 则直接跳到b的路径去寻找有没有bc
            while (cur.children[index] == null && cur != root) {
                cur = cur.fail;
            }
            cur = cur.children[index];
            // 如果当前路径以及所有失败路径都没有匹配到, 则把cur重置为root, 然后i++继续匹配
            if (cur == null) {
                cur = root;
            }
            TrieNode tmp = cur;
            while (tmp != root) {
                // 依次输出当前路径下所有有效的子字符串
                if (tmp.isEndingChar) {
                    // 计算当前子字符串的首字母坐标
                    int pos = i - tmp.length + 1;
                    System.out.println("匹配到字符串下标: " + pos + " 长度: " + tmp.length + " 字符串为: " + String.copyValueOf(text).substring(pos, pos + tmp.length));
                }
                tmp = tmp.fail;
            }
        }
    }

    // 纯Trie树: 查找所有存在trie树中的子串(没用AC自动机的多模实现, 为了与AC自动机相比较)
    public void matchWithoutFail(char[] text) {
        for (int i = 0; i < text.length; i++) {
            TrieNode cur = root;
            int index = text[i] - 'a';
            cur = cur.children[index];
            if (cur == null) {
                continue;
            }
            // AC自动机失败节点是跳到上一层, 如abc -> bc -> c
            // 因为这边没法直接跳到上一层, 所以是从前往后遍历, 如a -> ab -> abc
            for (int j = 1; j < text.length - i; j++) {
                if (cur.isEndingChar) {
                    // 当前子字符串的首字母坐标就是i
                    int pos = i;
                    System.out.println("匹配到字符串下标: " + pos + " 长度: " + cur.length + " 字符串为: " + String.copyValueOf(text).substring(pos, pos + cur.length));
                }
                // 寻找下一个子节点, 比如字符串abc 当前为ab, 则要寻找下一层的c节点, text[i + j]就表示从text中取出c字符串
                cur = cur.children[text[i + j] - 'a'];
                if (cur == null) {
                    break;
                }
            }
        }
    }

    public class TrieNode {
        public char data;
        public TrieNode[] children = new TrieNode[26];
        public boolean isEndingChar = false;
        // 当isEndingChar为true时, 记录字符串长度
        public int length = -1;
        // AC自动机: 失败指针
        public TrieNode fail;

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