Python中使用最小堆(Heap)数据结构实现Prim和Dijkstra算法

2023-04-11 00:00:00 算法数据结构最小

Prim算法的Python实现：

import heapq

def prim(adj_list):
    n = len(adj_list)
    visited = [False] * n
    heap = [(0, 0)]  # (distance, node)
    mst_cost = 0

    while len(heap) > 0:
        cost, node = heapq.heappop(heap)
        if visited[node]:
            continue
        visited[node] = True
        mst_cost += cost
        for neighbor, weight in adj_list[node]:
            heapq.heappush(heap, (weight, neighbor))

    return mst_cost

其中，adj_list是邻接表形式表示的图，如下所示：

adj_list = [[(1, 2), (2, 3)],
            [(0, 2), (2, 1), (3, 1)],
            [(0, 3), (1, 1), (3, 5), (4, 6)],
            [(1, 1), (2, 5), (4, 2)],
            [(2, 6), (3, 2)]]

Dijkstra算法的Python实现：

import heapq

def dijkstra(adj_list, start):
    n = len(adj_list)
    dist = [float('inf')] * n
    dist[start] = 0
    heap = [(0, start)]  # (distance, node)

    while len(heap) > 0:
        cost, node = heapq.heappop(heap)
        if cost > dist[node]:
            continue
        for neighbor, weight in adj_list[node]:
            new_cost = cost + weight
            if new_cost < dist[neighbor]:
                dist[neighbor] = new_cost
                heapq.heappush(heap, (new_cost, neighbor))

    return dist

其中，adj_list是邻接表形式表示的图，start是起点，如下所示：

adj_list = [[(1, 2), (2, 3)],
            [(0, 2), (2, 1), (3, 1)],
            [(0, 3), (1, 1), (3, 5), (4, 6)],
            [(1, 1), (2, 5), (4, 2)],
            [(2, 6), (3, 2)]]
start = 0

相关文章