最小生成树算法-- K算法、P算法

最小生成树

最小生成树 (Minimum Spanning Tree) ，假设一个图中有N个节点，用N-1条边将所有N个节点连通起来，要求这N-1条边的权值之和最小。常用的算法有：Kruskal（克鲁斯卡尔）算法和Prim（普里姆）算法。

Kruskal 算法

Kruskal 算法利用小根堆和并查集来实现，步骤：

1. 总是从权值最小的边开始考虑，依次考察权值依次变大的边
2. 当前的边要么进入最小生成树的集合，要么丢弃
3. 如果当前的边进入最小生成树的集合中，不会形成环，就要当前边，否则不要
4. 考察完所有边之后，最小生成树的集合就得到了

class UnionFind:
    """
    并查集基础算法
    """
    def __init__(self):
        self.parents = dict()
        self.size_map = dict()

    def makeSets(self, nodes: List[Node]):
        self.parents.clear()
        self.size_map.clear()

        for node in nodes:
            self.parents[node] = node
            self.size_map[node] = 1

    def findFather(self, cur: Node):
        """
        给你一个节点，请你往上到不能再往上，把代表节点返回
        :param cur:
        :return:
        """
        path_stack = []
        while cur != self.parents.get(cur):
            path_stack.append(cur)
            cur = self.parents.get(cur)

        while path_stack:
            self.parents[path_stack.pop()] = cur

        return cur

    def isSameSet(self, a, b):
        return self.findFather(a) == self.findFather(b)

    def union(self, a, b):
        a_head = self.findFather(a)
        b_head = self.findFather(b)

        if a_head != b_head:
            a_set_size = self.size_map.get(a_head)
            b_set_size = self.size_map.get(b_head)
            big = a_head if a_set_size >= b_set_size else b_head
            small = b_head if big == a_head else a_head
            self.parents[small] = big
            self.size_map[big] = a_set_size + b_set_size
            self.size_map.pop(small)

        return self.size_map.__len__()

    def getSetNum(self):
        return len(self.size_map)


class Kruskal:

    def kruskalMST(self, graph: Graph):
        unionfind = UnionFind()
        unionfind.makeSets(graph.nodes.values())

        # 从小的边到大的边，依次弹出，小根堆！
        hp = []
        for edge in graph.edges:
            heapq.heappush(hp, edge)

        res = set()
        while hp:
            edge = heapq.heappop(hp)
            if not unionfind.isSameSet(edge.src, edge.dst):
                res.add(edge)
                unionfind.union(edge.src, edge.dst)

        return res

Prim 算法

Prim算法使用两个容器：小根堆和集合。
小根堆用来存放所有被解锁的边，集合用来保存所有被解锁过的点。
每个边有四种状态：锁定、被解锁、选中、丢弃。

步骤：点->所有相邻的边，左右两端有不在点集之中的点的边中取最小权重->点->repeat

class Prim:
    def primMst(self, graph: Graph):
        # 解锁的边进入小根堆
        heap_edge = list()
        # 哪些点被解锁出来了
        nodeSet = set()
        # 被选中的边放在结果集合中
        result = set()

        for node in graph.nodes.values():  # 遍历每个点是避免有多个连通子图
            if node not in nodeSet:
                nodeSet.add(node)
                for edge in node.edges:
                    heapq.heappush(heap_edge, edge)

                while heap_edge:
                    edge = heapq.heappop(heap_edge)
                    dst = edge.dst
                    if dst not in nodeSet:
                    	result.add(edge)
                        nodeSet.add(dst)
                        for nextEdge in dst.edges:
                            heapq.heappush(heap_edge, nextEdge)
		return result