常用算法模板函数(Python)

并查集模板

# # ==> 并查集模板(附优化)
# author: [email protected]
class UnionFind():
    def __init__(self):
        self.roots = {}
        self.setCnt = 0   # 连通分量的个数
        # Union优化:存储根节点主导的集合的总节点数
        self.rootSizes = {}
    
    def add(self, x):
        if x not in self.roots:
            self.roots[x] = x
            self.rootSizes[x] = 1
            self.setCnt += 1
    
    def find(self, x):
        root = x
        while root != self.roots[root]:
            root = self.roots[root]
        # 优化:压缩路径
        while x != root:
            temp = self.roots[x]
            self.roots[x] = root
            x = temp
        return root
    
    def union(self, x, y):
        rootx, rooty = self.find(x), self.find(y)
        if rootx != rooty:
            # 优化:小树合并到大树上
            if self.rootSizes[rootx] < self.rootSizes[rooty]:
                self.roots[rootx] = rooty
                self.rootSizes[rooty] += self.rootSizes[rootx]
            else:
                self.roots[rooty] = rootx
                self.rootSizes[rootx] += self.rootSizes[rooty]
            self.setCnt -= 1

二分查找(左闭右开)

# # ==> 二分法(左闭右开)
# author: [email protected]
def lower_bound(nums,target):
    left,right=0,len(nums)
    while left<right:
        mid=(right-left)//2+left
        if nums[mid]<target:
            left=mid+1   # 维持left的绝对合理性;# left的左边一定小于target
        else:
            right=mid   # 逼出while循环
    return left

二叉树遍历

二叉树中序遍历和前序遍历(栈)
# # ==> 二叉树的中/前序遍历(栈)
# author: [email protected]
def inorder(root):
    stack=[]
    node=root
    while stack or node:
        while node:
            print(node.val) # 前序遍历
            stack.append(node)
            node=node.left
        node=stack.pop()
        print(node.val) # 中序遍历
        node=node.right
二叉树后序遍历(栈)
# # ==> 二叉树的后序遍历(栈)
# author: [email protected]
def afterOrder(root):
    result=[]
    stack=[]
    prev=None
    node=root
    while stack or node:
        while node:
            stack.append(node)
            node=node.left
        node=stack.pop()
        # 在中序遍历的基础上增加这部分。使用prev记录上一个访问的节点;如果node的右节点为空,那直接访问node节点即可,并更新prev和node,node更新为None是因为最终访问的是根节点,如果不置为空,将陷入死循环;如果不为空且右节点已经访问过(即等于prev),那么说明右节点已经全部访问了,那么访问node即可,并更新prev和node(同上);如果右节点不为空且没有访问过,则将node入栈,为了在右子树遍历完进行访问,并让node=node.right,让子树进行和node一样的操作
        if node.right is None or node.right==prev:
            result.append(node.val)
            prev=node
            node=None   # 方便下一次循环直接弹出当前节点的父节点,这个node会弹出两次,分为right==prev和right!=prev
        else:
            stack.append(node)  # 该节点还没处理,先处理右节点,先回到栈中
            node=node.right
    return result
二叉树的Morris算法前序遍历和中序遍历
# # ==> 二叉树的Morris算法前/中序遍历
# author: [email protected]
def preOrderByMorris(root) -> list[int]:
    # 第一步,特殊情况处理。树为空的情况下
    if root is None:
        return []
    # 第二步,遍历树直到当前结点为空
    cur = root
    while cur:
        if not cur.left:
            # 2.1.当前结点没有左子树,则更新当前结点为右结点(也只有这么走);并在更新之前访问当前结点
            print(cur.val)
            cur = cur.right
        else:
            # 2.2.当前结点有左子树的情况。找到左子树最右侧的结点rightMost;让rightMost的右指针线索指向当前结点并在更新后让当前结点继续往左走;如果rightMost已经指向了当前结点,说明处于回退阶段,左子树已经访问完成,断开rightMost的右线索并让当前结点往右走
            # 2.2.1.获取左子树的最右侧结点
            rightMost = cur.left
            while rightMost.right is not None and rightMost.right is not cur:
                rightMost = rightMost.right
            if rightMost.right is None:
                # 2.2.2.线索连接并更新当前结点往左走(前序遍历在构建线索时访问当前结点)
                print(cur.val)
                rightMost.right = cur
                cur = cur.left
            else:
                # 2.2.3.线索断开并更新当前结点往右走(中序遍历在断开线索时访问当前结点)
                # print(cur.val)
                rightMost.right = None
                cur = cur.right

字典树/前缀树模板

# # ==> 字典树模板
# author: [email protected]
class Trie:
    def __init__(self):
        self.sons={}     # 孩子节点
        self.end=False      # 是否是单词的尾节点
    
    # 往当前前缀树中插入单词
    def insert(self,word):
        cur=self
        for c in word:
            if c not in cur.sons:
                cur.sons[c]=Trie()
            cur=cur.sons[c]
        cur.end=True
    
    # 查找单词。返回:0->不存在该单词;1->存在word的路径,但是不存在单词;2->存在word的单词
    def find(self,word):
        cur=self
        for c in word:
            if c not in cur.sons:
                return 0
            cur=cur.sons[c]
        return 2 if cur.end else 1
    
    # 展现树;(非必要)
    def display(self):
        que=[self]
        print("->")
        while que:
            for i in range(len(que)):
                node=que.pop(0)
                for key,subNode in node.sons.items():
                    print(key,end=" ")
                    que.append(subNode)
            print("")
        print("=>")

辗转相除法求最大公约数

# # ==> 辗转相除法求最大公约数
# author: [email protected]
def gcd(m,n):
    maxVal,minVal=max(m,n),min(m,n)
    r=maxVal%minVal
    while r!=0:
        maxVal=minVal
        minVal=r
        r=maxVal%minVal
    return minVal

树状数组模板

树状数组(求区间和)
# ==> 树状数组(求区间和)
# author: [email protected]
class TreeArr():
    def __init__(self, length):
        self.length=length
        self.arr=[0]*self.length
    
    # 工具函数:二进制从右向左第一个1和其右边的0组成的数字
    def lowerbit(self,x):
        return x&(-x)   # 计算机的负数采用的是补码(和取反加1效果一致)
    
    # 基础函数:原数组arr[index]增加val,更新树状数组
    def add(self, index:int, val:int) -> None:
        while index < self.length:
            self.arr[index] += val
            index = index + self.lowerbit(index + 1)
    
    # 基础函数:arr[0]->arr[index]项前缀和
    def query(self, index:int) -> int:
        sum_ = 0
        while index >= 0:
            sum_ += self.arr[index]
            index -= self.lowerbit(index + 1)
        return sum_
    
    # 获取两个索引之间的和
    def sumRange(self, left: int, right: int) -> int:
        # 使用query的求前缀和函数进行实现
        pass
树状数组(求区间最大值)
# ==> 树状数组(求区间最大值)(需要维护原数组)
# author: [email protected]
class MaxTreeArr():
    def __init__(self, n:int):
        self.n = n
        self.nums = [-inf] * n
        self.arr = [-inf] * n
    
    def lowerbit(self, x:int) -> int:
        return x & (-x)
    
    # > 更新原数组的index下标处的元素值为max(nums[index],val),并更新树状数组中index的祖先结点
    def update(self, index:int, val:int) -> None:
        # 注意:这个地方对原数组的更新,一定要取最大值,否则可能导致覆盖最大值导致错误(跳过坑)
        self.nums[index] = max(self.nums[index], val)
        while index < self.n:
            self.arr[index] = max(self.arr[index], val)
            index += self.lowerbit(index + 1)
    
    # > 查找arr闭区间[left,right]之间的最大值
    def queryMax(self, left:int, right:int) -> int:
        ans = -inf
        index = right
        while index >= left:
            # l维护index维护的区间的左端点下标;如果l>=left,说明index维护的区间在[left,right]闭区间中,所以可以算最大值,反之,只能从原数组nums中选择,并将index自减1
            l = index - self.lowerbit(index + 1) + 1
            if l >= left:
                ans = max(ans, self.arr[index])
                index = l - 1
            else:
                ans = max(ans, self.nums[index])
                index -= 1
        return ans

线段树模板

求区间最大值的线段树模板
# ==> 线段树(维护区间最大值)
# author: [email protected]
class SegmentTreeToMax():
    def __init__(self, n:int):
        self.n = n  # 原数组的长度
        self.arr = [-inf] * (4 * n)

    # 基础函数:修改原数组中的index下标处的元素为val,并更新线段树
    def change(self, index:int, val:int, node:int, start:int, end:int) -> None:
        # 第一步,递归退出条件。到达index对应的叶结点
        if index == start and index == end:
            # > 防止后面的小元素替换大元素,所以加max(根据具体的修改场景可能会有修改)
            self.arr[node] = max(self.arr[node], val)
            return
        # 第二步,根据mid判断index是在当前node的左子树还是右子树上,并进行递归修改
        mid = (end - start) // 2 + start
        if index <= mid:
            # > index在node左子树的情况
            self.change(index, val, 2 * node + 1, start, mid)
        else:
            # > index在node的右子树的情况
            self.change(index, val, 2 * node + 2, mid + 1, end)
        # 第三步,更新当前结点,根据当前结点更新后的左右结点,更新当前结点的最大值
        self.arr[node] = max(self.arr[2 * node + 1], self.arr[2 * node + 2])
    
    # 基础函数:求区间范围内的最大值
    def rangeMax(self, left:int, right:int, node:int, start:int, end:int) -> int:
        # 第一步,递归退出条件。当node区间和[left,right]闭区间完全匹配时,递归退出
        if left == start and right == end:
            return self.arr[node]
        # 第二步,递归子问题
        mid = (end - start) // 2 + start
        if right <= mid:
            # > [left,right]区间完全在node结点的左子树区间上
            return self.rangeMax(left, right, 2 * node + 1, start, mid)
        elif left > mid:
            # > [left,right]区间完全在node结点的右子树区间上
            return self.rangeMax(left, right, 2 * node + 2, mid + 1, end)
        # > [left,right]区间部分在node的左子树上,部分在右子树上的情况
        return max(self.rangeMax(left, mid, 2 * node + 1, start, mid), self.rangeMax(mid + 1, right, 2 * node + 2, mid + 1, end))
    
    # 封装change
    def update(self, index:int, val:int) -> None:
        return self.change(index, val, 0, 0, self.n - 1)

    # 封装rangeMax
    def getRangeMax(self, left:int, right:int) -> int:
        # 注意:left>right时,会导致内存溢出
        return self.rangeMax(left, right, 0, 0, self.n - 1)
带懒标记的线段树(求区间和)
# ==> 带懒标记 求范围和 的线段树
# author: [email protected]
class SegmentTreeToSumWithLazy():
    def __init__(self, nums:list[int]):
        # 初始化线段树的存储数组并进行构建二叉平衡线段树(这里采用平衡二叉树,也可以用最大堆进行构建)
        self.n = len(nums)
        self.tree = [0] * (self.n * 4)   # 二叉平衡树的范围为4*n,如果使用最大堆的自底向上,则范围2*n即可;维护区间和
        self.lazy = [0] * (self.n * 4)  # 懒标记数组,lazy[node]表示node结点的子孙结点对应的区间待增加的值(不包括node结点)
        self.build(nums, 0, 0, self.n - 1)
    
    # 基础方法:构建树,在node树中,将nums[start:end+1]中的区间元素进行插入(node、start、end确定一个线段树结点)
    def build(self, nums:list[int], node:int, start:int, end:int) -> None:
        if start == end:
            self.tree[node] = nums[start]
            return 
        mid = (end - start) // 2 + start
        # 构建左子树
        self.build(nums, node * 2 + 1, start, mid)
        # 构建右子树
        self.build(nums, node * 2 + 2, mid + 1, end)
        # 和形式的线段树;tree[node]记录nums[start:end+1]之间元素的和
        self.tree[node] = self.tree[node * 2 + 1] + self.tree[node * 2 + 2]
    
    # 结点懒标记向下推送(基础函数):将node结点的懒标记推送到左右子结点中,并将自身的懒标记清空
    def pushDown(self, node:int, start:int, end:int) -> None:
        # 第一步,特殊判断。node结点的懒标记值为0,无需向下推送,直接退出
        if self.lazy[node] == 0:
            return 
        # 第二步,获取左右结点的结点编号和范围,更新左右结点的tree和lazy中的值
        mid = (end - start) // 2 + start
        leftChild = node * 2 + 1
        rightChild = node * 2 + 2
        self.tree[leftChild] += self.lazy[node] * (mid - start + 1)
        self.lazy[leftChild] += self.lazy[node]
        self.tree[rightChild] += self.lazy[node] * (end - mid)
        self.lazy[rightChild] += self.lazy[node]
        # 第三步,清空node对应的lazy值
        self.lazy[node] = 0

    # 区间修改(基础方法):在结点node对应的[start,end]区间中,将其中[left,right]区间部分的原数组值全部增加value
    def rangeUpdate(self, value:int, left:int, right:int, node:int, start:int, end:int) -> None:
        # 第一步,递归退出条件
        # 1.1.node结点区间[start,end]和目标区间[left,right]不存在交集时,直接退出
        if left > end or right < start:
            return 
        # 1.2.node结点区间[start,end]被目标区间[left,right]包含,直接更新lazy和tree数组。tree这种node结点对应值自增;在lazy数组中更新node结点的懒标记,lazy[node]自增value
        if left <= start and right >= end:
            self.tree[node] += value * (end - start + 1)
            self.lazy[node] += value
            return
        # 第二步,node区间和目标区间存在重叠的情况下,递归处理
        # 2.1.由于需要分到两个子区间中进行递归操作,所以node对应的lazy值需要向下推送到子结点中;本质就是将node结点的lazy值分配到左右结点中,更新tree和lazy中对应的值
        self.pushDown(node, start, end)
        # 2.2.递归范围更新两个子树
        leftChild, rightChild = node * 2 + 1, node * 2 + 2
        mid = (end - start) // 2 + start
        self.rangeUpdate(value, left, right, leftChild, start, mid)
        self.rangeUpdate(value, left, right, rightChild, mid + 1, end)
        # 第三步,更新当前结点的tree值(node中增加value的范围不确定,所以通过子结点来更新)
        self.tree[node] = self.tree[leftChild] + self.tree[rightChild]
    
    # 单点修改:调用区间函数rangeUpdate进行
    def pointUpdate(self, index:int, value:int, node:int, start:int, end:int) -> None:
        self.rangeUpdate(value, index, index, node, start, end)

    # 区间查询(基础方法):在结点node对应的[start,end]区间中,求[left,right]区间部分的原数组值的和
    def rangeSum(self, left:int, right:int, node:int, start:int, end:int) -> int:
        # 第一步,递归退出条件
        # 1.1.node结点区间和目标区间[left,right]不存在交集时,直接退出
        if left > end or right < start:
            return 0
        # 1.2.node结点区间被目标区间[left,right]包含,直接返回tree数组中node结点的值
        if left <= start and right >= end:
            return self.tree[node]
        # 第二步,node区间和目标区间存在重叠的情况下,递归处理
        # 2.1.由于需要分到两个子区间中进行递归操作,所以node对应的lazy值需要向下推送到子结点中;本质就是将node结点的lazy值分配到左右结点中,更新tree和lazy中对应的值
        self.pushDown(node, start, end)
        # 2.2.递归获取两个子树的范围和,相加进行返回
        mid = (end - start) // 2 + start
        return self.rangeSum(left, right, node * 2 + 1, start, mid) + self.rangeSum(left, right, node * 2 + 2, mid + 1, end)
带懒标记的线段树(求区间最大值)
# ==> 带懒标记 求范围最大值 的线段树【未经过题目测验,用了一个数组实验了一下】
# author: [email protected]
inf = float("inf")
class SegmentTreeToMaxWithLazy():
    def __init__(self, nums:list[int]):
        # 初始化线段树的存储数组并进行构建二叉平衡线段树(这里采用平衡二叉树,也可以用最大堆进行构建)
        self.n = len(nums)
        self.tree = [0] * (self.n * 4)   # 二叉平衡树的范围为4*n,如果使用最大堆的自底向上,则范围2*n即可;维护区间最大值
        self.lazy = [0] * (self.n * 4)  # 懒标记数组,lazy[node]表示node结点的子孙结点对应的区间待增加的值(不包括node结点)
        self.build(nums, 0, 0, self.n - 1)
    
    # 基础方法:构建树,在node树中,将nums[start:end+1]中的区间元素进行插入(node、start、end确定一个线段树结点)
    def build(self, nums:list[int], node:int, start:int, end:int) -> None:
        if start == end:
            self.tree[node] = nums[start]
            return 
        mid = (end - start) // 2 + start
        # 构建左子树
        self.build(nums, node * 2 + 1, start, mid)
        # 构建右子树
        self.build(nums, node * 2 + 2, mid + 1, end)
        # 和形式的线段树;tree[node]记录nums[start:end+1]之间元素的最大值
        self.tree[node] = max(self.tree[node * 2 + 1], self.tree[node * 2 + 2])
    
    # 结点懒标记向下推送(基础函数):将node结点的懒标记推送到左右子结点中,并将自身的懒标记清空
    def pushDown(self, node:int, start:int, end:int) -> None:
        # 第一步,特殊判断。node结点的懒标记值为0,无需向下推送,直接退出
        if self.lazy[node] == 0:
            return 
        # 第二步,获取左右结点的结点编号和范围,更新左右结点的tree和lazy中的值
        mid = (end - start) // 2 + start
        leftChild = node * 2 + 1
        rightChild = node * 2 + 2
        self.tree[leftChild] += self.lazy[node]
        self.lazy[leftChild] += self.lazy[node]
        self.tree[rightChild] += self.lazy[node]
        self.lazy[rightChild] += self.lazy[node]
        # 第三步,清空node对应的lazy值
        self.lazy[node] = 0

    # 区间修改(基础方法):在结点node对应的[start,end]区间中,将其中[left,right]区间部分的原数组值全部增加value
    def rangeUpdate(self, value:int, left:int, right:int, node:int, start:int, end:int) -> None:
        # 第一步,递归退出条件
        # 1.1.node结点区间[start,end]和目标区间[left,right]不存在交集时,直接退出
        if left > end or right < start:
            return 
        # 1.2.node结点区间[start,end]被目标区间[left,right]包含,直接更新lazy和tree数组。tree这种node结点对应值自增;在lazy数组中更新node结点的懒标记,lazy[node]自增value
        if left <= start and right >= end:
            self.tree[node] += value
            self.lazy[node] += value
            return
        # 第二步,node区间和目标区间存在重叠的情况下,递归处理
        # 2.1.由于需要分到两个子区间中进行递归操作,所以node对应的lazy值需要向下推送到子结点中;本质就是将node结点的lazy值分配到左右结点中,更新tree和lazy中对应的值
        self.pushDown(node, start, end)
        # 2.2.递归范围更新两个子树
        leftChild, rightChild = node * 2 + 1, node * 2 + 2
        mid = (end - start) // 2 + start
        self.rangeUpdate(value, left, right, leftChild, start, mid)
        self.rangeUpdate(value, left, right, rightChild, mid + 1, end)
        # 第三步,更新当前结点的tree值(node中增加value的范围不确定,所以通过子结点来更新)
        self.tree[node] = max(self.tree[leftChild], self.tree[rightChild])
    
    # 单点修改:调用区间函数rangeUpdate进行
    def pointUpdate(self, index:int, value:int, node:int, start:int, end:int) -> None:
        self.rangeUpdate(value, index, index, node, start, end)

    # 区间查询(基础方法):在结点node对应的[start,end]区间中,求[left,right]区间部分的原数组值的最大值
    def rangeMax(self, left:int, right:int, node:int, start:int, end:int) -> int:
        # 第一步,递归退出条件
        # 1.1.node结点区间和目标区间[left,right]不存在交集时,直接退出
        if left > end or right < start:
            return -inf
        # 1.2.node结点区间被目标区间[left,right]包含,直接返回tree数组中node结点的值
        if left <= start and right >= end:
            return self.tree[node]
        # 第二步,node区间和目标区间存在重叠的情况下,递归处理
        # 2.1.由于需要分到两个子区间中进行递归操作,所以node对应的lazy值需要向下推送到子结点中;本质就是将node结点的lazy值分配到左右结点中,更新tree和lazy中对应的值
        self.pushDown(node, start, end)
        # 2.2.递归获取两个子树的范围最大值,取最大值进行返回
        mid = (end - start) // 2 + start
        return max(self.rangeMax(left, right, node * 2 + 1, start, mid), self.rangeMax(left, right, node * 2 + 2, mid + 1, end))

拓扑排序算法

# author: [email protected]
from typing import Dict,List
from collections import deque
# ==> 通过List[List]结构的邻接表图和List形式的入度信息进行拓扑排序,返回拓扑排序序列,如果不存在,则返回空列表
def topoSort(adjListGraph:List[List[int]],inDegreeList:List[int]):
    que=deque()
    length=len(adjListGraph)
    for node in range(length):
        inDegree=inDegreeList[node]
        if inDegree==0:
            que.append(node)
    result=[]
    while que:
        node=que.popleft()
        result.append(node)
        for subNode in adjListGraph[node]:
            inDegreeList[subNode]-=1
            if inDegreeList[subNode]==0:
                que.append(subNode)
    return result if len(result)==len(adjListGraph) else []

最小生成树算法

prim算法
# author: [email protected]
import heapq
from typing import Dict,List
# ==> Prim算法模板:用于计算无向图的最小生成树及最小权值和
# graph:无向图的邻接表;item项例子:{节点:[[相邻边的权值,相邻边对面的节点],...],...}
# return:minWeightsSum为最小生成树的权值和;treeEdges为一个合法的最小生成树的边的列表(列表项:[节点,对面节点,两点之间的边的权值]);visited为最小生成树的节点,可以用来判断图中是否存在最小生成树
def primMinSpanningTree(graph:Dict[object,List[List]]):
    minWeightsSum,treeEdges=0,[]
    firstNode=list(graph.keys())[0] # *
    # 记录已经加入最小生成树的节点
    visited=set([firstNode])
    # 选择从firstNode开始,相邻的边加到堆中
    neighEdgesHeap=[item+[firstNode] for item in graph[firstNode]]
    heapq.heapify(neighEdgesHeap)
    while neighEdgesHeap:
        weight,node,startNode=heapq.heappop(neighEdgesHeap) # node2为node的weight权值对应的边的对面的节点
        if node in visited:    # 这个地方必须进行判断,否则会造成重复添加已访问节点,造成最小权值和偏大(因为前面遍历的节点可能将未遍历的共同相邻节点重复添加到堆中)
            continue
        minWeightsSum+=weight
        treeEdges.append([startNode,node,weight])
        visited.add(node)
        # 遍历新访问的节点的边,加入堆中
        for nextWeight,nextNode in graph[node]:
            if nextNode not in visited:
                heapq.heappush(neighEdgesHeap,[nextWeight,nextNode,node])
    return minWeightsSum,treeEdges,visited

单源最短路径算法

dijkstra算法
# author: [email protected]
# ==> dijkstra求单源最短路径
# params:graph为邻接表,邻接表结构为graph[FromNode_i]=[[Dist1,ToNode1],[Dist2,ToNode2],...],三维数组;startNode为源节点
# return:dists为各个节点的最短路径的长度;pathsPrevs为各个节点的最短路径的前驱节点
# waring: *处修改可以改成求最大路径
# 举例:获取最后一个节点的最短路径
# paths=[]
# node=length-1
# while node!=-1:
    # paths.append(node)
    # node=pathsPrevs[node]
from typing import List
import heapq
inf=float("inf")
def dijkstraMinDist(graph:List[List[List]],startNode:int):
    length=len(graph)
    # 第一步,构建距离数组、前驱节点数组、距离最小堆。并初始化startNode的最小距离为0
    distances=[inf]*length  # *各个节点到startNode的最小距离
    distances[startNode]=0  # *初始化startNode到startNode的最小距离0
    pathsPrevs=[-1]*length      # 最短路径的最后节点的前驱节点
    distsHeap=[[0,startNode]]   # *距离优先队列,项的结构为[距离startNode的距离,node]
    # 第二步,遍历。每次从最小堆中弹出一个(直到堆为空);在弹出的节点距离小于或等于distances数组中的对应节点的距离的情况下,遍历节点的子节点,如果到达子节点的距离小于原先距离,更新距离数组、前驱节点数组、距离最小堆
    while distsHeap:
        dist,node=heapq.heappop(distsHeap)
        if dist>distances[node]:    # *排除同一个
            continue
        for edgeWeight,subNode in graph[node]:
            thisDist=edgeWeight+distances[node]
            if thisDist<distances[subNode]: # *
                distances[subNode]=thisDist
                pathsPrevs[subNode]=node
                heapq.heappush(distsHeap,[thisDist,subNode])    # *
    return distances,pathsPrevs
spfa算法
# author: [email protected]
from collections import deque
# ==> 求单源最短路径,可以包含负权边(不能有负环)
# graph:[(结点,有向边权重)]
def spfa(graph:list[list[int]], start:int, end:int):
    inf = float("inf")
    n = len(graph)
    # 第一步,构建维护变量。que队列用来维护当前层的结点;visited用来记录各个结点当前是否在队列中,如果当前结点已经在队列中了,那么就不用重复入队了;dists数组用来维护各个结点到起始点的最短距离,初始化为inf,start结点初始化为0;paths用来记录最短路径上各个结点的前驱结点,初始化为None,start结点初始化为-1;cnt用来维护每个结点入队的次数,用来验证是否存在负环,如果存在某个结点入队次数超过n,那么说明存在负环
    que = deque([start])
    visited = [False] * n
    visited[start] = True
    dists = [inf] * n
    dists[start] = 0
    paths = [None] * n
    paths[start] = -1
    cnt = [0] * n
    cnt[0] = 1
    # 第二步,循环从队列中取出结点,类似广度优先搜索。一层层的从队列中弹出元素,并遍历相邻结点;先根据相邻结点的新最短距离来胡决定更新更近的相邻结点;相邻结点距离更近的情况下将没有在当前队列中的相邻结点添加到队列中,继续进行遍历(两个if的关键点),并更新入队次数和判断是否为负环
    while que:
        for _ in range(len(que)):
            node = que.popleft()
            visited[node] = False
            for neighNode, weight in graph[node]:
                if dists[neighNode] > dists[node] + weight:
                    dists[neighNode] = dists[node] + weight
                    paths[neighNode] = node
                    if not visited[neighNode]:
                        que.append(neighNode)
                        visited[neighNode] = True
                        cnt[neighNode] += 1
                        if cnt[neighNode] > n:
                            raise Exception("存在负环,请检查图是否有错误!")
    # print(dists)
    # print(paths)
    return dists, paths

二分图最大匹配数算法

匈牙利算法
# author: [email protected]
# ==> 匈牙利算法(时间复杂度:O(V*E))
# params:graph:二分图的邻接表;leftNodes:二分图左集合的所有结点数组
# return:返回二分图中的最大匹配数(等于最小点覆盖数)
from collections import defaultdict
def hungarian(graph:dict[int,list[int]], leftNodes:list[int]) -> int:
    # 第一步,构建维护变量。rightMatch维护右集合中各个结点的匹配状态
    rightMatch = {}
    # 第二步,构建dfs递归函数。递归任务:判断左集合中结点u是否能成功匹配。如果u相邻结点v没有匹配,则直接匹配;如果v已经匹配,则将v标记为已访问,dfs(match[v])观察是否能让match[v]结点匹配非v结点来释放v结点,让u进行匹配,即贪心地进行匹配。
    def dfs(u:int) -> bool:
        # 2.1.枚举u的相邻的右集合结点v
        for v in graph[u]:
            # 2.2.判断v是否已访问,已访问则跳过
            if not visited[v]:
                # 2.3.标记v为已访问,并根据条件判断是否能匹配,能匹配的情况下记录右结点v的匹配状态到rightMatch中
                visited[v] = True
                if v not in rightMatch or dfs(rightMatch[v]):
                    rightMatch[v] = u
                    return True
        # 2.4.如果最终都没能匹配一个,则返回False
        return False
    # 第三步,枚举每一个左集合结点u,调用dfs判断各个u结点是否能成功进行匹配,并记录最大匹配数到ans中
    ans = 0
    for u in leftNodes:
        visited = defaultdict(bool)
        if dfs(u):
            ans += 1
    return ans
Hopcroft-karp算法
# author: [email protected]
# ==> Hopcroft-karp算法(时间复杂度:O(sqrt(V)*E))
# params:graph:二分图的邻接表;leftNodes:二分图左集合的所有结点数组;
# return:返回二分图中的最大匹配数(等于最小点覆盖数)
from collections import deque
def hopcroftKarp(graph:list[list[int]], leftNodes:list[int]):
    inf = float("inf")
    # 第一步,构建维护变量。match_维护二分图中各个匹配关系;dist维护各个交替路中距离左集合中未匹配点的最短距离,dist中映射值为inf代表结点已经匹配或者不存在以该结点开头的增广路径
    match_ = {}
    dist = {}
    # 第二步,构建广搜函数。bfs功能:返回在match_和dist的情况下,是否还存在更多的增广路径,同时更新dist
    def bfs():
        # 2.1.将左侧集合的未匹配结点全部加入到队列中,作为BFS的起点;同时将左集合中已经匹配的结点的距离值设置为inf
        queue = deque()
        for u in leftNodes:
            if u not in match_:
                dist[u] = 0
                queue.append(u)
            else:
                dist[u] = inf
        # 2.2.广搜二分图。使用found维护图中是否还存在增广路径
        found = False
        while queue:
            u = queue.popleft()
            for v in graph[u]:
                if v not in match_:
                    # 2.2.1.如果右侧结点v没有在match_中,说明v结点还没有匹配,表示图中还存在增广路径,设置found=True
                    found = True
                elif dist[match_[v]] == inf:
                    # 2.2.2.如果右侧结点v在match_中且对应的match_[v]结点还没有被广搜到(即dist[match_[v]==inf]),说明v结点已经匹配,则将v的匹配结点match_[v]加入到广搜队列中,继续寻找增广路径;并更新match_[v]结点所在的最小广搜层次
                    queue.append(match_[v])
                    dist[match_[v]] = dist[u] + 1
        # 2.3.返回图中是否还存在增广路径
        return found
    # 第三步,构建受限深搜函数。深搜任务:判断以结点u开头是否存在增广路径;如果存在,则在该路径上进行结点匹配,反转该增广路径。
    def dfs(u):
        for v in graph.get(u, []):
            if v not in match_ or (dist.get(match_[v], inf) == dist[u] + 1 and dfs(match_[v])):
                # 3.1.如果结点v没有匹配,则u->v就是最短的增广路径,直接匹配u和v并返回即可;如果结点v已经匹配了,只有满足match_[v]是广搜中结点u的下一层级时(即dist[match_[v]]==dist[u]+1)(这个条件也能过滤掉递归时match_[v]遍历到相邻结点v的情况),才能继续dfs判断是否存在增广路径,如果存在则匹配并反转该增广路径
                match_[v] = u
                match_[u] = v
                return True
        # 3.2.如果不存在以u开头的增广路径,则在dist中进行标记,并返回False
        dist[u] = inf
        return False
    # 第四步,循环地进行bfs判断图中是否存在增广路径;如果存在则遍历所有左侧没有匹配的结点u,dfs寻找是否存在u开头的最短增广路径,并在dfs中更新match_和dist,如果存在以u开头的增广路径,则通过反转增广路径进行匹配,并将ans自增1
    ans = 0
    while bfs():
        for u in leftNodes:
            if u not in match_:
                if dfs(u):
                    ans += 1
    return ans

懒删除小根堆模板

# author: [email protected]
# ==> 模板:懒删除小根堆
from collections import defaultdict
from heapq import heappush, heappop, heappushpop
class LazyMinHeap:
    def __init__(self):
        # 第一步,构建维护变量。heap维护最小堆;size维护最小堆中元素的个数;sum维护堆中所有元素的和;cnts维护堆中各个有效元素的频数;delayDelete维护待删除的各个元素的个数;
        self.heap = []
        self.size = 0
        self.sum = 0
        self.cnts = defaultdict(int)
        self.delayDelete = defaultdict(int)
    
    # 修剪堆顶部待删除的元素
    def prune(self):
        while self.heap:
            topVal = self.heap[0]
            if self.delayDelete[topVal] > 0:
                heappop(self.heap)
                self.delayDelete[topVal] -= 1
            else:
                break
    
    # 懒删除堆中的元素;这里必须保证val在堆中
    def delete(self, val:int):
        if self.cnts[val] > 0:
            self.cnts[val] -= 1
            self.delayDelete[val] += 1
            self.size -= 1
            self.sum -= val
            return True
        return False

    # 判断堆是否为空
    def empty(self):
        self.prune()
        return len(self.heap) == 0

    # 获取顶部元素
    def top(self):
        if self.empty():
            return
        return self.heap[0]

    # 堆中添加元素
    def push(self, val:int):
        if self.delayDelete[val] > 0:
            self.delayDelete[val] -= 1
        else:
            heappush(self.heap, val)
        self.cnts[val] += 1
        self.sum += val
        self.size += 1
    
    # 堆顶弹出元素
    def pop(self):
        if self.empty():
            return
        val = heappop(self.heap)
        self.cnts[val] -= 1
        self.size -= 1
        self.sum -= val
        return val
    
    def __str__(self):
        return f"size={self.size}\nsum={self.sum}\ndelayDelete={self.delayDelete}\ncnts={self.cnts}\nheap={self.heap}\n"

你可能感兴趣的:(python,算法,模板函数)