Python中如何实现后缀树字符串匹配算法

class Node:
    def __init__(self, start, end):
        self.children = {}  # 子节点
        self.start = start  # 起始位置
        self.end = end      # 结束位置

def build_suffix_tree(s):
    root = Node(-1, -1)  # 创建根节点
    for i in range(len(s)):
        cur = root
        for j in range(i, len(s)):
            if s[j] not in cur.children:
                cur.children[s[j]] = Node(j, len(s)-1)
                break
            else:   # 子节点中有当前字符
                child = cur.children[s[j]]
                k = child.start    # 子节点起始位置
                while k <= child.end and j < len(s) and s[j] == s[k]:
                    j += 1
                    k += 1
                if k <= child.end:
                    # 将现有节点分裂成两个，一个是相同的前缀，另一个是后缀
                    new_node = Node(child.start, k-1)
                    child.start = k
                    new_node.children[s[k]] = child
                    cur.children[s[j-1]] = new_node  # 挂在新的节点上
                cur = child
    return root

def match_suffix_tree(s, root, pattern):
    cur = root
    for i in range(len(pattern)):
        if pattern[i] not in cur.children:
            return False
        else:
            child = cur.children[pattern[i]]
            k = child.start
            j = i
            while k <= child.end and j < len(pattern) and pattern[j] == s[k]:
                j += 1
                k += 1
            if k <= child.end:
                return False
            cur = child
    return True

# 测试
s = "pidancode.com"
root = build_suffix_tree(s)
print(match_suffix_tree(s, root, "pi"))     # True
print(match_suffix_tree(s, root, "pida"))   # True
print(match_suffix_tree(s, root, "code."))  # True
print(match_suffix_tree(s, root, "python")) # False