Pythonでクラスを用いたグラフの実装と幅優先探索

はじめに

最近、競プロを始めました。現在AtCoderで茶色です。

競プロの勉強をする中で、必ず出てくるのがグラフ（ネットワーク）の実装です。
よくある実装としては隣接リスト表現といって、ノードの番号を入れ子のリストに格納する方法があります。
ただ僕にとっては、その実装は直感的ではないと感じました。せっかくオブジェクト指向の言語を使っているのだから、グラフもクラスを使って表現したいと思ったのです。

そこで、僕がいつもやっているクラスを用いたグラフの実装方法について書いてみたいと思います。言語はPythonです。

よくあるグラフの実装：隣接リスト

たとえば、こんなグラフを表現したいとします。

隣接リスト表現を使うと、こんな感じになるかと思います。

G = [
    [1], # 頂点0の隣接頂点は1
    [0, 2, 3], # 頂点1の隣接頂点は0, 2, 3
    [1], # 頂点2の隣接頂点は1
    [1], # 頂点3の隣接頂点は1
]

クラスを用いたグラフの実装

次に、クラスを用いた実装をしてみます。まず、頂点を表すためにこんなNodeクラスを用意します。

class Node:
    def __init__(self) -> None:
        self.neighbors: list[Node] = [] # 隣接頂点のリストをアトリビュートとして持つ

そして、上記のグラフはこのように表現できます。

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

これで、各頂点オブジェクトのアトリビュートを見るだけで、隣接頂点が分かるようになっています。

幅優先探索の実装

これを使って、幅優先探索を実装してみます。
例題として、書籍『競技プログラミングの鉄則』にある以下の問題を解いてみます。

$N$頂点$M$辺の無向グラフが与えられます。各頂点には$1$から$N$までの番号が付けられており、$i$番目の辺は頂点$A_i$と頂点$B_i$を結んでいます。$1$以上$N$以下の全ての整数$k$について、以下の問いに答えてください。
頂点$1$から頂点$k$まで、辺を何本かたどって移動することを考えるとき、たどるべき辺の本数の最小値を出力せよ。ただし、移動不可能な場合は$−1$を出力せよ。

まず、頂点を表すNodeクラスを用意します。頂点0からの距離もこのクラスのアトリビュートとして持たせます。

class Node:
    def __init__(self) -> None:
        self.neighbors: list[Node] = []
        self.distance = -1  # 頂点0からの距離。ただし-1で初期化

続いて、標準入力を読み取ってグラフを構成します。

# 標準入力から頂点の数Nと辺の数Mを取得
N, M = map(int, input().split())

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

# 標準入力からM個の辺の情報を取得してグラフを構成
for _ in range(M):
    # 結ぶ頂点の番号a, bを取得
    a, b = map(int, input().split())
    # 頂点リストは0番から始まっているので、-1してインデックスを調整する
    a -= 1
    b -= 1
    
    # 頂点aの隣接頂点リストに頂点bを、頂点bの隣接頂点リストに頂点aを加える
    nodes[a].neighbors.append(nodes[b])
    nodes[b].neighbors.append(nodes[a])

そして、幅優先探索の実装は以下のようになります。

# キューの定義
d: deque[Node] = deque()

# キューに0番目の頂点を追加
d.append(nodes[0])

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

# 幅優先探索：キューが空になるまで以下の処理を行う
while d:
    # キューから頂点を取得
    node = d.popleft()
    # 各隣接頂点について処理を行う
    for neighbor in node.neighbors:
        # 距離が未確定（-1）であれば距離を確定し、キューに入れる
        if neighbor.distance == -1:
            neighbor.distance = node.distance + 1
            d.append(neighbor)

# 出力
for node in nodes:
    print(node.distance)

隣接リスト表現を使う場合は、隣接リスト自体に距離の情報を持たせることができないので、距離を格納するための変数を別途用意する必要があります。
一方この実装だと、距離の情報も頂点オブジェクト自身に持たせることができます。隣接頂点情報の取得も含めて、より直感的なコードになっているのではと思っています。

参考：隣接リスト表現による実装例（『競技プログラミングの鉄則』サポートページより）
https://github.com/E869120/kyopro-tessoku/blob/main/codes/python/chap09/answer_A63.py

おわりに

競技プログラミングはその場限りで動くコードを書けばよいので、わざわざクラスを定義する人は少数派だと思います。ただ個人的には、クラスを定義することにより考えを整理でき、結果的に早く問題を解けることも多いように感じています。
だれかの参考になれば幸いです。