Dijkstra迪杰斯特拉算法Python版本

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Time : 2022/1/31 14:14
# @Author : Jin Echo
# @File : Dijkstra.py

import heapq
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import random


def dijkstra(adj, s):
    """

    :param adj: n*n的邻接矩阵
    :param s: 起点下标
    :return:
    """
    n = adj.shape[0]  # 顶点数
    pqueue = []
    heapq.heappush(pqueue, (0, s))  # <距离 起点>二元组
    visited = set()
    pre = np.arange(n)
    # 初始化
    distance = np.ones((n,)) * np.inf
    distance[s] = 0

    cnt = 0
    while len(visited) != n and len(pqueue):
        cnt += 1
        min_dis, min_node = heapq.heappop(pqueue)
        if min_node in visited:
            continue
        else:
            visited.add(min_node)
        neighbors = np.where((adj[min_node, :] < 1) & (adj[min_node, :] > 0))[0]  # 提取当前点的邻接点
        for neighbor in neighbors:
            if neighbor not in visited:
                if min_dis + adj[min_node][neighbor] < distance[neighbor]:
                    distance[neighbor] = min_dis + adj[min_node][neighbor]
                    heapq.heappush(pqueue, (distance[neighbor], neighbor))
                    pre[neighbor] = min_node
                    
    print(cnt)
    return pre, distance


def get_sequence(pre, node):
    """
    返回起点到指定节点的最短路径
    :param pre:
    :param node:
    """
    res = [node]
    while pre[node] != node:
        node = pre[node]
        res.append(node)
    res.reverse()
    return res


if __name__ == '__main__':
    np.random.seed(1)
    random.seed(1)
    n, m = 20, 4
    G = nx.barabasi_albert_graph(n, m)  # 生成图结构
    plt.figure(1)
    nx.draw(G, with_labels=True, font_weight='bold')

    edges_idx = np.array(G.edges)
    # print(edges_idx)
    num_edges = (n - m) * m  # 图中的边数，也等于edges_idx.shape[0]
    assert num_edges == edges_idx.shape[0]

    # 赋值边权
    adj = np.ones([n] * 2) * np.inf
    # adj[range(n), range(n)] = 0  # 主对角线为0
    adj[edges_idx[:, 0], edges_idx[:, 1]] = np.random.rand(num_edges)
    adj[edges_idx[:, 1], edges_idx[:, 0]] = adj[edges_idx[:, 0], edges_idx[:, 1]]

    s = 0
    pre, distance = dijkstra(adj, s)
    print(pre)
    print(distance)
    print(get_sequence(pre, n - 1))
    plt.show()