Python 堆优化Prim MLE
查看原帖
Python 堆优化Prim MLE
1077867
SwordRain楼主2024/4/2 23:55

样例过了,全MLE,不知是否是Python自身的缘故,我把代码放在这,留给大佬来调

import sys
import heapq
input = lambda: sys.stdin.readline().strip()
def prim(graph, n, h):
    pqueue = []
    vis = [False] * (n + 1)
    for node, weight in graph[1]:
        if h[0] >= h[node - 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
            if i == n:
                break
            for node, weight in graph[vertex]:
                if h[vertex - 1] >= h[node - 1]:
                    heapq.heappush(pqueue, (weight, node))
    print(i, ans)
def main():
    n, m = map(int, input().split())
    h = list(map(int, input().split()))
    graph = {i: [] for i in range(1, n + 1)}
    for _ in range(m):
        start, end, weight = map(int, input().split())
        graph[start].append((end, weight))
    prim(graph, n, h)
main()
2024/4/2 23:55
加载中...