Leetcode刷题笔记1 图论part07

卡码网 53 寻宝

prim算法

prim算法核心就是三步，称为prim三部曲：

1. 第一步，选距离生成树最近节点
2. 第二步，最近节点加入生成树
3. 第三步，更新非生成树节点到生成树的距离（即更新minDist数组）

def prim(v, e, edges):
    import sys
    import heapq
    grid = [[10001] * (v + 1) for _ in range(v + 1)]
 
    for edge in edges:
        x, y, k = edge
        grid[x][y] = k
        grid[y][x] = k
    minDist = [10001] * (v + 1)
    isInTree = [False] * (v + 1)
 
    for i in range(1, v):
        cur = -1
        minVal = sys.maxsize
 
        for j in range(1, v + 1):
            if not isInTree[j] and minDist[j] < minVal:
                minVal = minDist[j] 
                cur = j
 
        isInTree[cur] = True
 
        for j in range(1, v + 1):
            if not isInTree[j] and grid[cur][j] < minDist[j]:
                minDist[j] = grid[cur][j]
 
    result = sum(minDist[2:v + 1])
    return result
 
if __name__ == '__main__':
    import sys
    input = sys.stdin.read
    data = input().split()
    v = int(data[0])
    e = int(data[1])
    index = 2
    edges = []
    for i in range(e):
        x = int(data[index])
        y = int(data[index + 1])
        k = int(data[index + 2])
        edges.append((x, y, k))
        index += 3
    ans = prim(v, e, edges)
    print(ans)

 kruskal 算法

kruscal的思路：

• 边的权值排序，因为要优先选最小的边加入到生成树里
• 遍历排序后的边
  • 如果边首尾的两个节点在同一个集合，说明如果连上这条边图中会出现环
  • 如果边首尾的两个节点不在同一个集合，加入到最小生成树，并把两节点加入同一个集合

class Edge:
    def __init__(self, l, r, val):
        self.l = l
        self.r = r
        self.val = val
n = 10001
father = list(range(n))
def init():
    global father 
    father = list(range(n))
def find(u):
    if u != father[u]:
        father[u] = find(father[u])
    return father[u]
def join(u, v):
    u = find(u)
    v = find(v)
    if u != v:
        father[v] = u
 
def kruskal(v, edges):
    edges.sort(key = lambda edge: edge.val)
    init()
    result = 0
    for edge in edges:
        x = find(edge.l)
        y = find(edge.r)
        if x != y:
            result += edge.val
            join(x, y)
    return result
 
if __name__ == '__main__':
    import sys
    input = sys.stdin.read
    data = input().split()
    v = int(data[0])
    e = int(data[1])
    index = 2
    edges = []
    for i in range(e):
        x = int(data[index])
        y = int(data[index + 1])
        k = int(data[index + 2])
        edges.append(Edge(x, y, k))
        index += 3
    ans = kruskal(v, edges)
    print(ans)