ペナルティが最小となるグラフ内の2頂点間のパスを見つけるPythonプログラム

問題概要

無向の重み付きグラフが与えられ、ノード a からノード b へ至る「最小ペナルティ」のパスを求めることを考えます。ここでいうパスのペナルティとは、そのパスに含まれるすべての辺の重みをビット単位OR（論理和）で結合した値のことです。つまり、ペナルティが最小となるパスを見つけ出し、もし2つのノード間にパスが存在しない場合は -1 を返す必要があります。

具体例

たとえば、次のようなグラフが与えられたとします。

始点 s = 1、終点 e = 3 のとき、出力は 15 になります。

頂点1と頂点3の間には2つのパスが存在します。最適なパスは 1 → 2 → 3 であり、このパスのコストは (10 OR 5) = 15 となります。

解法の考え方

この問題は、ダイクストラ法を応用することで効率的に解けます。ポイントは、ビット単位ORは値を重ねるほど減ることがない（単調非減少）という性質です。このため、「現時点でのペナルティが小さい状態から順に探索する」という優先度付きキュー（ヒープ）を使った手法が有効に機能します。

具体的には、以下の手順に従います。

• 補助関数 helper() を定義します。引数は G（隣接リスト形式のグラフ）、s（始点）、e（終点）です。
  • v := 訪問済みの (コスト, 頂点) の組を記録する新しい集合
  • c := サイズ n の新しいリスト。全要素を無限大（inf）で初期化
  • heap := ペア (0, s) を格納した新しいヒープ
  • ヒープが空になるまで、以下を繰り返します
    • ヒープから最小の項目を取り出し、コスト cst と現在の頂点 cur に分解する
    • c[cur] := cst と c[cur] のうち小さい方
    • (cst, cur) がすでに v に存在する場合は、次の反復へ進む
    • cur が e と一致した場合は、c[cur] を返す
    • (cst, cur) を v に追加する
    • G[cur] 内の各隣接頂点 neighbor とその辺の重み n_cost について、((n_cost OR cst), neighbor) をヒープにプッシュする
  • c[e] を返す
• G := n + 1 個の空リストからなる新しいリストを作成する
• edges の各要素について
  • u := item[0]、v := item[1]、w := item[2]
  • G[u] の末尾にペア (v, w) を追加する
  • G[v] の末尾にペア (u, w) を追加する（無向グラフのため両方向に登録）
• ans := helper(G, s, e)
• ans が inf と等しい場合は -1 を返し、それ以外の場合は ans を返す

実装例

理解を深めるために、以下の実装を見てみましょう。

import heapq
from math import inf

def helper(G, s, e):
    v = set()
    c = [inf] * len(G)
    heap = [(0, s)]
    while len(heap) > 0:
        cst, cur = heapq.heappop(heap)
        c[cur] = min(cst, c[cur])
        if (cst, cur) in v:
            continue
        if cur == e:
            return c[cur]
        v.add((cst, cur))
        for neighbor, n_cost in G[cur]:
            heapq.heappush(heap, (n_cost | cst, neighbor))
    return c[e]

def solve(n, edges, s, e):
    G = [[] for _ in range(n + 1)]
    for item in edges:
        u, v, w = map(int, item)
        G[u].append((v, w))
        G[v].append((u, w))
    ans = helper(G, s, e)
    return -1 if ans == inf else ans

print(solve(4, [(1, 2, 10), (2, 3, 5), (2, 4, 15), (1, 4, 20)], 1, 3))

入力

4, [(1, 2, 10), (2, 3, 5), (2, 4, 15), (1, 4, 20)], 1, 3

出力

15