图论算法——Dijistra和Floyd算法求最短路径

数模图论部分练习时用python实现了一个能调用Dijistra和Floyd算法的最短路径类。可以直接复制使用。

# FLoyd 和 Dijkstra算法

import numpy as np
from copy import deepcopy

class allpair_shortestpath:
    """
    Parameters:
    @__ajmatrix: the adjacent matrix with respect to the graph to be manipulated
    @__nodes_dict: a dict that mapping node index to its name
    @class path_and_nodes: a member class as shortest paths representation
    """
    class path_and_nodes:
        def __init__(self, dist, next_nodes, kind):
            self.dist, self.next_nodes, self.kind = (dist, next_nodes, kind)
            
    res = path_and_nodes(None, None, '')
    
    def __init__(self, ajmatrix, nodes_dict):
        self.__ajmatrix, self.__nodes_dict = (ajmatrix, nodes_dict)
        self.__nodes_count = ajmatrix.shape[0]
        self.__edges_count = np.sum(np.bitwise_and(ajmatrix>=0, ajmatrix<np.inf))
    def floyd(self):
        dist = deepcopy(self.__ajmatrix)
        next_nodes = np.hstack([np.full((self.__nodes_count,1), i) for i in range(self.__nodes_count)])
        
        for k in range(self.__nodes_count):
            for i in range(self.__nodes_count):
                for j in range(self.__nodes_count):
                    if dist[i][j] > dist[i][k] + dist[k][j]:
                        dist[i][j] = dist[i][k] + dist[k][j]
                        next_nodes[i][j] = next_nodes[i][k]
        next_nodes[np.bitwise_or(dist <0, dist>=np.inf)] = -1
        self.res.dist, self.res.next_nodes = (dist, next_nodes)
        return self.path_and_nodes(dist, next_nodes, 'Floyd')
    
    def floyd_route(self, pos):
        u, v = pos
        if self.res.next_nodes[u][v] < 0:
            return ""
        path = []
        path.append(self.__nodes_dict[u])
        while u != v:
            u = self.res.next_nodes[u][v]
            path.append(self.__nodes_dict[u])
        return '=>'.join(path)
    
    def dijkstra(self, source, dest=None):
        mask = np.bitwise_and(self.__ajmatrix>0, self.__ajmatrix<np.inf)
        dist = deepcopy(self.__ajmatrix[source,:])
        pred = np.full(dist.shape, -1)
        pred[mask[source,:]] = source
        unvisited = np.full(dist.shape, True)
#         print(unvisited)
#         print(np.bitwise_not(unvisited))

        while np.sum(unvisited) != 0:
            distcopy = deepcopy(dist)
#             print(distcopy)
#             print(unvisited, np.invert(unvisited))
            distcopy[np.invert(unvisited)] = np.inf
#             print(distcopy)
            i = np.argmin(distcopy)
            if np.count_nonzero(unvisited) == 1:
                i = np.flatnonzero(unvisited)
            unvisited[i] = False
#             print(i)
            for j in np.flatnonzero(mask[i,:]):
                if dist[j] > dist[i] + self.__ajmatrix[i,j]:
                    dist[j] = dist[i] + self.__ajmatrix[i,j]
                    pred[j] = i
        self.res.dist, self.res.next_nodes, self.res.kind = (dist, pred, 'Dijkstra')
        return self.path_and_nodes(dist, pred, 'Dijkstra')
    
    def dijkstra_route(self, pair):
        u, v = pair
#         self.dijkstra(u, v)
        assert self.res.kind == 'Dijkstra'
        path = []
        if self.res.next_nodes[v] >= 0 or u == v:
            while v >= 0:
                path.insert(0, self.__nodes_dict[v])
                v = self.res.next_nodes[v]
        return '=>'.join(path)

if __name__ == '__main__':
    ajmatrix = np.array([
        [0, 6, 3, 1, -1, -1, -1, -1, -1],
        [-1, 0, -1, -1, 1, -1, -1, -1, -1],
        [-1, 2, 0, 2, -1, -1, -1, -1, -1],
        [-1, -1, -1, 0, -1, 10, -1, -1, -1],
        [-1, -1, -1, 6, 0, 4, 3, 6, -1],
        [-1, -1, -1, -1, 10, 0, 2, -1, -1],
        [-1, -1, -1, -1, -1, -1, 0, 4, -1],
        [-1, -1, -1, -1, -1, -1, -1, 0, -1],
        [-1, -1, -1, -1, 2, -1, -1, 3, 0]
    ], dtype=np.float)
    ajmatrix[ajmatrix <0] = np.inf
    nodes_dict = dict(zip(range(ajmatrix.shape[0]), ['v'+str(i+1) for i in range(ajmatrix.shape[0])]))
    
    obj = allpair_shortestpath(ajmatrix, nodes_dict)
    # print(obj._allpair_shortestpath__ajmatrix)
    res = obj.floyd()
    print(res.dist, res.next_nodes, sep='\n')
    print(obj.floyd_route((0,7)))

    obj = allpair_shortestpath(ajmatrix, nodes_dict)
    # print(obj._allpair_shortestpath__ajmatrix)
    res = obj.dijkstra(0)
    print(res.dist, res.next_nodes, sep='\n')
    print(obj.dijkstra_route((0,7)))