【代码随想录算法训练营——Day57（Day56周日休息）】图论——53.寻宝

卡码网题目链接

https://kamacoder.com/problempage.php?pid=1053

题解

53.寻宝

本题是模板题, 代码是仿照题解里的c++代码写出来的. prim和kruskal都是.

代码

#53.寻宝
#prim法
if __name__ == "__main__":
    v, e = map(int, input().split())
    graph = [[10001] * (v + 1) for _ in range(v + 1)]
    for _ in range(e):
        x, y, k = map(int, input().split())
        graph[x][y] = k
        graph[y][x] = k
    minDist = [10001] * (v + 1)
    isInTree = [False] * (v + 1)
    for i in range(1, v): # n - 1条边连接n个结点
        cur = -1
        minVal = float('inf')
        for j in range(1, v + 1):
            if isInTree[j] == False and minDist[j] < minVal:
                cur = j
                minVal = minDist[j]
        isInTree[cur] = True
        for j in range(1, v + 1):
            if isInTree[j] == False and graph[cur][j] < minDist[j]:
                minDist[j] = graph[cur][j]
    result = 0
    for i in range(2, v + 1):
        result += minDist[i]
    print(result)

#kruskal法
class Edge:
    def __init__(self, l, r, val):
        self.l = l
        self.r = r
        self.val = val

n = 10001
father = [-1] * n

def init():
    for i in range(n):
        father[i] = i
def find(u):
    if father[u] == u:
        return u
    else:
        father[u] = find(father[u])
        return father[u]
def isSame(u, v):
    u = find(u)
    v = find(v)
    return u == v
def join(u, v):
    u = find(u)
    v = find(v)
    if u == v:
        return
    father[v] = u

if __name__ == "__main__":
    v, e = map(int, input().split())
    result = 0
    edges = []
    for _ in range(e):
        v1, v2, val = map(int, input().split())
        edges.append(Edge(v1, v2, val))
    edges.sort(key = lambda x: x.val)
    init()
    for edge in edges:
        x = find(edge.l)
        y = find(edge.r)
        if x != y:
            result += edge.val
            join(x, y)
    print(result)