python实现AC自动机

ac自动机可以看成带指针的字典树，每个节点的指针指向了当前节点的最大后缀的位置。在建立字典树后，可以层次遍历字典树来构建fail指针，根节点的直接孩子（第一层节点）的fail指针肯定是指向根节点的，之后的节点需要看其父节点的fail指针指向的节点，如果父节点的fail指针指向的节点的孩子中存在当前节点，则将当前节点的fail指针指向该节点，如果不存在，需要沿着fail指针一直向上寻找直至根节点，如果根节点的孩子中也找不到，说明树中不存在该字符，将当前节点的fail指针指向根节点就行了。最后，在查找时也类似，如果当前节点的孩子中查不到，也需要沿着fail指针行走直至根节点。举个例子，关键字[“i”, “is”,“ssippi”]建好树并构建好fail指针后如下图所示：

上图中，右边的第一个i，父节点的fail指针指向第一个s，但该s的孩子中并不存在i，所以需要沿着fail指针向上查找，发现根节点root下存在i，因此fail指针指向这个i。另外，若节点代表了关键词的结尾，则标记了该词的长度，如果该节点的fail指针指向的节点也有长度，则将长度合并至该节点（见上图右下角的i），以便于最后按照长度查找关键词。以下是python实现的代码：

from queue import Queue
from typing import List, Dict, Iterable


class Node(object):

    def __init__(self, char: str):
        self.char = char  # 节点代表的字符
        self.children = {
     }  # 节点的孩子，键为字符，值为节点对象
        self.fail = None  # fail指针，root的指针为None
        self.exist = []  # 如果节点为单词结尾，存放单词的长度


class AhoCorasick(object):
    """AC自动机"""
    def __init__(self, keywords: Iterable[str] = None):
        self.root = Node("root")
        self.finalized = False
        if keywords is not None:
            for keyword in set(keywords):
                self.add(keyword)

    def add(self, keyword: str):
        if self.finalized:
            raise RuntimeError('Tree has been finalized!')
        node = self.root
        for char in keyword:
            if char not in node.children:
                node.children[char] = Node(char)
            node = node.children[char]
        node.exist.append(len(keyword))

    def contains(self, keyword: str) -> bool:
        node = self.root
        for char in keyword:
            if char not in node.children:
                return False
            node = node.children[char]
        return True

    def finalize(self):
        """构建fail指针"""
        queue = Queue()
        queue.put(self.root)
        # 对树进行层次遍历
        while not queue.empty():
            node = queue.get()
            for char in node.children:
                child = node.children[char]
                f_node = node.fail
                # 关键点！需要沿着fail指针向上追溯直至根节点
                while f_node is not None:
                    if char in f_node.children:
                        # 如果该指针指向的节点的孩子中有该字符，则字符节点的fail指针需指向它
                        f_child = f_node.children[char]
                        child.fail = f_child
                        # 同时将长度合并过来，以便最后输出
                        if f_child.exist:
                            child.exist.extend(f_child.exist)
                        break
                    f_node = f_node.fail
                # 如果到根节点也没找到，则将fail指针指向根节点
                if f_node is None:
                    child.fail = self.root
                queue.put(child)
        self.finalized = True

    def search_in(self, text: str) -> Dict[str, List[int]]:
        """在一段文本中查找关键字及其开始位置（可能重复多个）"""
        result = dict()
        if not self.finalized:
            self.finalize()
        node = self.root
        for i, char in enumerate(text):
            matched = True
            # 如果当前节点的孩子中找不到该字符
            while char not in node.children:
                # fail指针为None，说明走到了根节点，找不到匹配的
                if node.fail is None:
                    matched = False
                    break
                # 将fail指针指向的节点作为当前节点
                node = node.fail
            if matched:
                # 找到匹配，将匹配到的孩子节点作为当前节点
                node = node.children[char]
                if node.exist:
                    # 如果该节点存在多个长度，则输出多个关键词
                    for length in node.exist:
                        start = i - length + 1
                        word = text[start: start + length]
                        if word not in result:
                            result[word] = []
                        result[word].append(start)
        return result

AC自动机的一个好处是，可以在大文本中快速的查找多个关键词，下面是几个测试案例：

---

输入：
keywords = [“i”, “is”,“ssippi”]
text = ‘misissippi’
输出：
{‘i’: [1, 3, 6, 9], ‘is’: [1, 3], ‘ssippi’: [4]}

---

输入：
keywords = [“what”, “hat”, “ver”, “er”]
text = ‘whatever, err … , wherever’
输出：
{‘what’: [0], ‘hat’: [1], ‘ver’: [5, 25], ‘er’: [6, 10, 22, 26]}

---

输入：
keywords = [“江苏省”, “苏州市”, “姑苏区”, “吴中区”]
text = ‘阿斯蒂芬江苏省，苏州市，阿斯蒂芬吴中区，苏州市’
输出：
{‘江苏省’: [4], ‘苏州市’: [8, 20], ‘吴中区’: [16]}