[算法入土之路]前缀树

前缀树:

        用处:

                    查询一个字符串是否被添加在一个序列中
                    查询序列中是否有一个以...字符为前缀的字符串
                    以上时间复杂度为常数级 

数据结构

class TreeNode:
    '''
    :param self.pass: 序列经过标识 如果一个序列经过了该节点 则此值+1
            如果该节点为根节点 则 pass 代表了 该树有多少个序列
    :param self.end:  序列结束标识  如果一个序列的最后一个元素是此节点 则此值+1
            如果该节点为根节点 则 end 代表了 该树有多少个空序列
    :param self.next: 子节点集
    '''

    def __init__(self, pass_=0, end=0):
        self.pass_ = pass_
        self.end = end
        self.next = {}  # 子节点集

实现

class PrefixTree:
    def __init__(self, root=None):
        '''
        接收传入的根节点 如果没有传则创建一个新的单一节点作为根节点
        :param root: 根节点
        '''
        self.root = root if root else TreeNode(pass_=0)

    def insert(self, seq=None):
        '''
        插入一个序列
        :param seq: 元素序列
        :return:
        '''
        if not seq:
            self.root.pass_ += 1
            self.root.end += 1
            return
        node = self.root  # 从根节点开始
        for elem in seq:  # 遍历序列中的每一个元素
            node.pass_ += 1  # 碰到的节点 pass+1 代表一个元素经过了该节点
            if elem not in node.next:  # 如果本元素未在此节点的后代中, 则新建一个节点
                cur_node = TreeNode()
                node.next[elem] = cur_node  # 将新建的节点加入到当层节点的后代中
            node = node.next[elem]  # 进入下一层节点
        node.pass_ += 1
        node.end += 1  # 最后一个元素的节点的序列结束标识 + 1

    def search(self, seq=None):
        '''
        查询序列被添加过几次
        :param seq: 序列
        :return: 添加次数
        '''
        if not seq:  # 如果为空序列
            return self.root.end
        node = self.root
        # 遍历序列, 找到seq[-1]指向的节点
        for elem in seq:
            if elem in node.next:
                node = node.next[elem]
            else:  # 如果没有找到 则说明没有 seq 序列
                return 0
        return node.end

    def start_with_seq_index(self, seq=None):
        '''
        以序列 seq 为开头的的序列有多少个
        :param seq:
        :return: 满足条件的序列个数
        '''
        if not seq:  # 如果为空 则返回树中的序列个数
            return self.root.pass_
        node = self.root
        for elem in seq:  # 遍历 seq 序列
            if elem in node.next:  # 如果 elem 在当前节点的子集中
                node = node.next[elem]  # 向下传递
            else:  # 如果没找到 则表示没有以seq为开头的序列
                return 0
        # 找到seq[-1]对应的节点  返回node.pass_ 即为所求
        return node.pass_

    def remove(self, seq=None):
        if not seq:  # 如果是空序列
            if self.root.pass_ > len(self.root.next):  # 如果树中含有空序列
                # 删除空序列
                self.root.pass_ -= 1
                self.root.end -= 1
            return  # 返回
        if self.search(seq):  # 先查看下树中是否存在该序列
            self.root.pass_ -= 1  # 树中序列个数 -1
            node = self.root
            for elem in seq:  # 遍历序列元素
                node.next[elem].pass_ -= 1  # 途径节点的经过标识 -1
                if not node.next.get(elem).pass_:  # 如果下个节点的 经过标识为0
                    # 则表示要删除的序列为1 直接从根节点的子集中删除该序列即可
                    del node.next[elem]
                    return
                else:
                    node = node.next[elem]  # 向下传递
            node.end -= 1  # node 为 seq[-1] 所指向的节点 此操作代表以 seq[-1] 为结尾的序列数目-1

打印前缀树

def print_dic(dic, inner_index=0):
    '''
    打印前缀树
    :param dic: 根节点的子节点集
    :param inner_index: # 第几层
    :return:
    '''
    for key, value in dic.items():
        print(key, value.pass_, value.end, inner_index, end=" | ")
        if not (inner_index + 1) % 8:
            print()
            print("".rjust(79, "-"))
        if value:
            print_dic(value.next, inner_index + 1)
        else:
            print()

测试主函数

if __name__ == '__main__':
    t = PrefixTree()
    t.insert("zhangsan")
    t.insert("zhangsaa")
    # t.insert("zhangs")
    # t.remove("zhangsan")
    t.insert()

    print_dic(t.root.next)
    # res = t.search("zhangs")
    res = t.start_with_seq_index("z")
    print(f"res: {res}")