后缀自动机的python实现

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2019-02-25 14:32:50
# @Author  : Sheldon (thisisscret@qq.com)
# @Blog    : 谢耳朵的派森笔记
# @Link    : https://www.cnblogs.com/shld/
# @Version : 0.0.1

from collections import OrderedDict
import logging
from copy import deepcopy

logger = logging.getLogger("sam")
logger.propagate = False
logger.handlers.clear()
logger.setLevel(level=logging.INFO)
# handler = logging.FileHandler("log.txt")
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)


class State:
    def __init__(self, length=None, link=None, trans=None):
        self.length = length
        self.link = link
        self.trans = trans


class Sam:
    def __init__(self, s=""):
        self.last = 0
        self.sz = 1
        self.st = OrderedDict({0: State(0, -1, OrderedDict())})
        self.samstr = ""
        self(s)

    @staticmethod
    def _dict2str(dct):
        return ' '.join(f'{k}->{w}' for k, w in dct.items())

    def _extend(self, c):
        cur = self.sz

        self.st[cur] = State(None, None, OrderedDict())
        self.st[cur].length = self.st[self.last].length + 1
        logger.debug(f"状态{cur}添加到后缀自动机")
        logger.debug(f"上一个节点状态是{self.last}")
        logger.debug(f"状态{cur}的长度是{self.st[self.last].length + 1}")

        # 如果后缀自动机最近转移里面没有当前字符，则添加该字符，并将状态指向当前状态
        # 继续沿着后缀连接走，进行上述操作直到到达第一个状态或者转移中有此字符
        p = self.last

        while p > -1 and self.st[p].trans.get(c) is None:
            logger.debug(
                f"状态{p}的转移:{self._dict2str(self.st[p].trans)}{('','不')[self.st[p].trans.get(c) is None]}包含字符{c}")

            self.st[p].trans[c] = cur
            logger.debug(f"把{c}添加进状态{p}的转移")
            logger.debug(f"开始查找状态{p}的后缀链接...")
            p = self.st[p].link
            logger.debug(f"后缀链接为状态{p}")

        # 如果后缀链接走到底了，没有相同的，则后缀链接指向0状态，即空字符串
        if p == -1:
            self.st[cur].link = 0
            logger.debug(f"状态{cur}沿后缀链接走到了初始状态,将其的后缀链接指向状态0")

        # 如果找到上一状态的转移里有c字符,找到转移c的另一状态
        else:
            q = self.st[p].trans[c]
            logger.debug(f"在状态{p}的转移中找到了字符{c}，{c}指向状态{q}")
            # 如果q状态与p状态相连，则当前状态的后缀链接指向q状态
            if self.st[p].length + 1 == self.st[q].length:
                self.st[cur].link = q
                logger.debug(f"状态{p}的长度比状态{q}少1，把当前状态{cur}的后缀链接指向状态{q}")
            # 如果不相连则开一个新状态,长度为p状态的下一个状态，后缀链接与转移指向q
            # 搜索状态p，若c转移为q，则指向新状态，并搜索后缀链接的状态重复指向新状态
            #   直到状态转移不为q，跳出
            else:
                self.sz += 1
                logger.debug(f"状态{p}的长度与状态{q}的长度不连续，新建状态{self.sz}")
                self.st[self.sz] = State(self.st[p].length + 1, self.st[q].link, deepcopy(self.st[q].trans))
                logger.debug(f"新状态{self.sz}的长度为状态{p}的长度加1")
                while p > -1 and self.st[p].trans.get(c) == q:
                    logger.debug(f"状态{p}转移中{c}的转移是状态{q},把指向{q}的转移改为指向新状态{self.sz}")
                    self.st[p].trans[c] = self.sz
                    logger.debug(f"查找状态{p}的后缀链接...")
                    p = self.st[p].link
                    logger.debug(f"后缀链接为状态{p}")
                # 把当前状态与q的后缀链接指向新状态
                self.st[q].link = self.st[cur].link = self.sz
                logger.debug(f"把状态{q}与当前状态{cur}的后缀链接指向新状态{self.sz}")
        logger.debug(f"当前状态{cur}的长度为:{self.st[cur].length}")
        logger.debug(f"当前状态{cur}的后缀链接为:{self.st[cur].link}")
        logger.debug(f"当前状态{cur}的转移为:{self._dict2str(self.st[cur].trans)}")
        logger.debug(f"当前状态{cur}建立完成")
        logger.debug("\n")
        # 状态索引占位
        self.sz += 1
        self.last = cur
        self.samstr += c

    def __add__(self, addition):
        new_sam = deepcopy(self)
        if isinstance(addition, str):
            for c in addition:
                new_sam._extend(c)
        elif isinstance(addition, Sam):
            for c in addition.samstr:
                new_sam._extend(c)
        return new_sam

    def __call__(self, s):
        for c in s:
            self._extend(c)

    def __str__(self):
        pts = ""
        for k, w in self.st.items():
            pts += f"状态{k}: 长度:{w.length}  后缀链接-->{w.link}  转移:{self._dict2str(w.trans)}\n"
        return pts

    def __getitem__(self, key):
        return self.st[key]

if __name__ == "__main__":
    
    s1 = "aabbabd"
    s2 = "rtrbWHWUKkkh"
    sam1 = Sam(s1)
    sam2 = Sam(s2)
    
    print(sam1)
    
    #加的时候不会改变原后缀自动机
    sam3 = sam1+sam2
    sam4 = sam1+s2
    
    #调用则在原自动机上更新
    sam1(s2)

• 调试的时候日志级别设为DEBUG就行了
• 有空的时候更新相应字符串处理的算法