Python代码求调
查看原帖
Python代码求调
1077867
SwordRain楼主2024/4/3 00:35

堆优化Prim,全WA

import sys
import heapq
input = lambda: sys.stdin.readline().strip()  # 快读
def prim(graph, n):
    pqueue = []  # 空优先队列
    vis = [False] * (n + 1)  # 记录是否走过该点
    for node, weight in graph[1]:  # 将1节点的所有边加入到优先队列
        heapq.heappush(pqueue, (weight, node))
    i = 1  # 记录节点数量
    ans = 0  # 记录总长度
    while pqueue:  # 当优先队列不为空
        weight, vertex = heapq.heappop(pqueue)  # 获取优先级最高的距离和节点
        if not vis[vertex]:  # 如果没见过
            vis[vertex] = True  # 设置为见过
            ans += weight  # 加上这一段距离
            i += 1  # 节点数+1
            if i == n:  # 如果所有节点都走过
                break  # 直接跳出循环
            for node, weight in graph[vertex]:  # 把该节点的所有边加入到优先队列
                heapq.heappush(pqueue, (weight, node))
    print(i, ans)  # 打印答案
def main():
    n, m = map(int, input().split())  # 输入n,m
    h = list(map(int, input().split()))  # 获取每个景点的高度
    graph = {i: [] for i in range(1, n + 1)}  # 邻接表
    for _ in range(m):  # 循环m次
        start, end, weight = map(int, input().split())  # 读取起点,终点,距离
        if h[start - 1] >= h[end - 1]:  # 如果start的高度大于等于end的高度
            graph[start].append((end, weight))  # 把end和weight加入到start
        if h[start - 1] <= h[end - 1]:  # 如果end的高度大于等于start的高度
            graph[end].append((start, weight))  # 把start和weight加入到end
    prim(graph, n)  # 运行prim算法
main()
2024/4/3 00:35
加载中...