Pythonでクラスを用いた重み付きグラフとダイクストラ法の実装

はじめに

こちらの記事の続編です。

上記の記事では、クラスを用いたグラフの実装を紹介しました。
この記事では、同じくクラスを用いて、重み付きグラフの実装を紹介します。ついでにダイクストラ法の実装も紹介します。

よくある重み付きグラフの実装

たとえば、こんなグラフを表現したいとします。頂点0と頂点1を繋ぐ辺の重み（距離）が3、という具合です。

よくある実装としては、隣接リスト表現の拡張で、隣接ノード番号の代わりに隣接ノード番号とそれを繋ぐ辺の重みのタプルを格納するものです。

G = [
    [(1, 3)], # 頂点0の隣接頂点は1で距離は3
    [(0, 3), (2, 1), (3, 2)], # 頂点1の隣接頂点は0, 2, 3で、距離はそれぞれ3, 1, 2
    [(1, 1)], # 頂点2の隣接頂点は1で距離は1
    [(1, 2)], # 頂点3の隣接頂点は1で距離は2
]

これでも全ての情報は表現できているのですが、頂点番号や重みをタプルから取得するためには添字を用いてアクセスする必要があり、あまり直感的ではないと感じています。

クラスを用いた重み付きグラフの実装

続いてクラスで実装するとどうなるかを見ていきます。まず、下記のようなNodeクラスとLinkクラスを用意します。

class Node:
    def __init__(self) -> None:
        self.links: list[Link] = [] # 辺のリスト

class Link:
    def __init__(self, distance: int, to_node: Node):
        self.distance = distance # 辺の重み
        self.to_node = to_node # 辺が繋がる先のノード

頂点クラスは、自分に繋がっている辺のリストを属性として持ちます。辺クラスは繋がる先の頂点を属性として持っちます。

これを用いると、上記のグラフはこのように表現できます。

nodes = [Node() for _ in range(4)]
nodes[0].links = [Link(3, nodes[1])] # 頂点0の隣接頂点は1で距離は3
nodes[1].links = [
    Link(3, nodes[0]), # 頂点1の隣接頂点は0で距離は3
    Link(1, nodes[2]), # 頂点1の隣接頂点は2で距離は1
    Link(2, nodes[3])  # 頂点1の隣接頂点は3で距離は2
] # # 頂点1の隣接頂点は0, 2, 3
nodes[2].links = [Link(1, nodes[1])] # 頂点2の隣接頂点は1で距離は1
nodes[3].links = [Link(2, nodes[1])] # 頂点3の隣接頂点は1で距離は2

ダイクストラ法の実装

これを使ってダイクストラ法を実装してみます。
例題として、書籍『競技プログラミングの鉄則』にある以下の問題を解いてみます。

重み付き無向グラフに対する最短経路問題を解いてください。 具体的には、以下のようなグラフが与えられるとき、頂点$1$から各頂点までの最短経路長を求めてください。

• 頂点数は $N$ 、辺数は$M$である
• $i$番目の辺は頂点$A_i$と頂点B_iを結び、長さはC_iである

まず、頂点を表すNodeクラスと辺を表すLinkクラスを用意します。
頂点クラスの属性には、頂点0からの距離と、距離が確定したかを表すフラグを持たせます。

class Node:
    def __init__(self) -> None:
        self.links: list[Link] = []
        self.distance = sys.maxisize # 頂点0からの距離（十分大きな数で初期化）
        self.finalized = False # 距離が確定したかを表すフラグ

class Link:
    def __init__(self, distance: int, to_node: Node):
        self.distance = distance
        self.to_node = to_node

標準入力を読み取り、頂点と辺を構成します。

# 頂点数N、辺数M
N, M = map(int, input().split())

# N個の頂点を用意
nodes: list[Node] = [Node() for _ in range(N)]

# M個の辺を用意
for _ in range(M):
    # 頂点Aと頂点Bを結ぶ、距離Cの辺
    A, B, C = map(int, input().split())
    # 頂点A,Bを取得（頂点リストのインデックスは0始まりなので-1して調整）
    nodeA = nodes[A - 1]
    nodeB = nodes[B - 1]
    # 頂点A,Bに辺を追加
    nodeA.links.append(Link(C, nodeB))
    nodeB.links.append(Link(C, nodeA))

続いて、ダイクストラ法では優先度付きキューを用いますが、そのキューに入れる要素を作るためのクラスを用意します。
キューに頂点そのものを入れることはできません。キューに入る2つ以上の要素が同じ頂点を指し示すケースもあるためです。

class QueueElement:
    def __init__(self, node: Node):
        # 要素が指し示す頂点
        self.node = node
        # 要素が指し示す頂点の距離
        # 頂点クラスのdistance属性は次々と暫定的に書き換わるが、このキュー要素を作った時点での暫定距離を保存する
        self.distance = node.distance

    # 優先度付きキュー内で用いるため、大小比較をするためのメソッド
    def __lt__(self, other):
        return self.distance < other.distance

以上の道具を使って、ダイクストラ法を実装します。

# 優先度付きキューを用意
q: list[QueueElement] = []

# 頂点0の距離を0で初期化
nodes[0].distance = 0

# キューに頂点0を入れる
heapq.heappush(q, QueueElement(nodes[0]))

# ダイクストラ法
while q:
    # 優先度付きキューから暫定距離が最小の頂点を取り出す
    queue_element = heapq.heappop(q)
    node = queue_element.node

    # 確定済みの頂点であれば処理をスキップし、確定前であれば確定させる
    if node.finalized:
        continue
    node.finalized = True

    # 頂点の各隣接頂点について、距離の暫定値を更新しキューに入れる
    for link in node.links:
        distance = node.distance + link.distance
        to_node = link.to_node
        if distance < to_node.distance:
            to_node.distance = distance
            heapq.heappush(q, QueueElement(to_node))

# 出力
for node in nodes:
    if node.distance == sys.maxsize:
        print(-1)
    else:
        print(node.distance)

おわりに

個人的には競プロでクラスを使うのはオススメなんですが、少数派なのかな…