[Python] ダイクストラ法

Python で、ダイクストラ法を使って、重み付きの有向グラフの単一始点最短経路問題を解きます。

以下の続きです。

ダイクストラ法

ダイクストラ法（だいくすとらほう、英: Dijkstra’s algorithm）はグラフ理論における辺の重みが非負数の場合の単一始点最短経路問題を解くための最良優先探索によるアルゴリズムである。

出典: フリー百科事典『ウィキペディア（Wikipedia）』

距離の∞は、 sys.maxsize を用いて表します。

sys.maxsize

頂点の集合の優先度付きキューは、heapq を用いて実装します。

heapq — ヒープキューアルゴリズム

優先度付きキューは以下で扱っています。

以下のようなグラフを探索します。

a を出発地点、e を目的地とした場合、最短経路は「a -> c -> b -> e」、移動距離は7になるはずです。

import heapq
import sys

class Vertex(object):
    def __init__(self, id):
        self.id = id
        self.adjacent = {}
        # 開始点からの距離
        self.distance = sys.maxsize
        self.visited = False
        # 前の点を記録し、最短経路を出力する時に使う
        self.previous = None

    def add_neighbor(self, neighbor, weight=0):
        self.adjacent[neighbor] = weight
    
    def get_neighbors(self):
        return self.adjacent

    def get_connections(self):
        return self.adjacent.keys()  

    def get_vertex_id(self):
        return self.id

    def get_weight(self, neighbor):
        return self.adjacent[neighbor]
    
    def set_distance(self, distance):
        self.distance = distance

    def get_distance(self):
        return self.distance

    def set_visited(self):
        self.visited= True

    def set_previous(self, previous_vertex):
        self.previous = previous_vertex 
    
    def get_previous(self):
        return self.previous

    # heapq を使うために必要。定義しないと以下のようなエラーが出る。
    # TypeError: '<' not supported between instances of 'Vertex' and 'Vertex'
    def __lt__(self, other):
        return self.distance < other.distance


class Graph(object):
    def __init__(self):
        self.vertex_dict = {}
        self.num_vertex = 0
    
    def __iter__(self):
        return iter(self.vertex_dict.values())

    def add_vertex(self, id):
        self.num_vertex = self.num_vertex + 1
        new_vertex = Vertex(id)
        self.vertex_dict[id] = new_vertex
        return new_vertex

    def get_vertex(self, id):
        if id in self.vertex_dict:
            return self.vertex_dict[id]
        else:
            return None

    def add_edge(self, frm, to, weight=0):
        if frm not in self.vertex_dict:
            self.add_vertex(frm)
        if to not in self.vertex_dict:
            self.add_vertex(to)
        self.vertex_dict[frm].add_neighbor(self.vertex_dict[to], weight)

    def get_vertices(self):
        return self.vertex_dict.keys()

    def get_edges(self):
        edges = []
        for v in self.vertex_dict.values():
            for w in v.get_connections():
                vid = v.get_vertex_id()
                wid = w.get_vertex_id()
                edges.append((vid, wid, v.get_weight(w)))
        return edges


def dijkstra(graph, source, destination):
    # ソースの距離を0に設定
    source.set_distance(0)

    # (距離, 頂点)のタプルのリストを作り、heapq で優先度付きキューにする。
    priority_queue = [(vertex.get_distance(), vertex) for vertex in graph]
    heapq.heapify(priority_queue)

    # while (Q が空集合ではない )
    while len(priority_queue):
        # Qから距離が最小である頂点を取り出す
        nearest_vrtex = heapq.heappop(priority_queue)
        current = nearest_vrtex[1]
        current.set_visited()

        #  for each ( u からの辺がある各頂点 )
        for next in current.adjacent:
            if next.visited:
                continue
            # alt = d(u)+length(u,v)
            new_distance = current.get_distance() + current.get_weight(next)

            # if ( d(v) > alt )
            if next.get_distance() > new_distance:
                next.set_distance(new_distance)
                next.set_previous(current)
                heapq.heappush(priority_queue, (next.get_distance(), next))
            
    for vertex in graph.vertex_dict.values():
        print(f'{source.get_vertex_id()} から {vertex.get_vertex_id()} への移動距離: {vertex.get_distance()}')

    shortest_path = [destination.get_vertex_id()]
    while destination.previous:
        shortest_path.append(destination.previous.get_vertex_id())
        destination = destination.previous
    print(f'{source.get_vertex_id()} から {vertex.get_vertex_id()} への最短移動経路:', shortest_path[::-1])
    

if __name__ == '__main__':

    graph = Graph()
    graph.add_vertex('a')
    graph.add_vertex('b')
    graph.add_vertex('c')
    graph.add_vertex('d')
    graph.add_vertex('e')
    graph.add_edge('a', 'b', 4)  
    graph.add_edge('a', 'c', 1)
    graph.add_edge('b', 'e', 4)
    graph.add_edge('c', 'b', 2)
    graph.add_edge('c', 'd', 4)
    graph.add_edge('d', 'e', 4)

    source = graph.get_vertex('a')
    destination = graph.get_vertex('e')
    dijkstra(graph, source, destination)
    # a から a への移動距離: 0
    # a から b への移動距離: 3
    # a から c への移動距離: 1
    # a から d への移動距離: 5
    # a から e への移動距離: 7
    # a から e への最短移動経路: ['a', 'c', 'b', 'e']

想定通りの結果を得ることができました。