Pythonで最初のノードから最後のノードまでの制限付きパスの数を求めるプログラム

無向の重み付き連結グラフがあるとします。グラフは n 個のノードを持ち、それぞれのノードには 1 から n までのラベルが付けられています。始点から終点へのパスとは [z0, z1, z2, ..., zk] のようなノードの列のことで、z0 が始点ノード、zk が終点ノードであり、隣り合うノード zi と zi+1 の間（0 ≤ i ≤ k-1）には必ず辺が存在します。パスの距離は、そのパスが通る辺の重みの総和として定義されます。また、dist(x) は「ノード n からノード x までの最短距離」を表すものとします。制限付きパス（restricted path）とは、すべての i（0 ≤ i ≤ k-1）について dist(zi) > dist(zi+1) という条件を満たす特別なパスのことです。ここでの課題は、ノード 1 からノード n までの制限付きパスの本数を求めることです。答えが非常に大きくなる可能性があるため、10^9 + 7 で割った余りを返してください。

例として、次のようなグラフが入力された場合を考えます。

このとき出力は 3 になります。制限付きパスは (1,2,5)、(1,2,3,5)、(1,3,5) の 3 本が存在するためです。

解き方のアプローチ

この問題は、ダイクストラ法と動的計画法（DP）を組み合わせることで効率的に解けます。以下の手順に従います。

• 与えられた辺リストをもとに、グラフの隣接リストを作成する
• サイズ (n+1) の配列 paths を用意し、すべて 0 で初期化する
• paths[n] := 1 と設定する（終点自身へのパスは 1 通り）
• サイズ (n+1) の配列 dists を用意し、すべて -1 で初期化する（未確定の印）
• 優先度付きキュー q を用意し、最初に (0, n) を挿入する
• q が空になるまで、次の処理を繰り返す
  • (dist, node) := キューの先頭要素を取り出す
  • dists[node] が -1 以外（すでに確定済み）なら、次の反復へスキップする
  • dists[node] := dist として距離を確定させる
  • node に隣接する各ノード v と辺の重み w について、
    • dists[v] が -1 の場合：(dist + w, v) をキューに挿入する
    • dists[v] < dists[node] の場合：paths[node] := paths[node] + paths[v] として経路数を累積する
  • node が 1 と等しければ、paths[node] mod (10^9 + 7) を返す
• ループが終了したら 0 を返す

なぜこのアルゴリズムが機能するのか

終点であるノード n を起点にダイクストラ法を実行すると、ノードは「ノード n からの距離が近い順」に確定していきます。したがって、あるノードが確定した時点で、そのノードより dist が小さい隣接ノードは必ず先に確定済みになっており、paths[v] の値はすでに最終的な値になっています。この性質により、「dist(zi) > dist(zi+1)」という制限を満たすパスの数を、paths[node] = Σ paths[v]（v は dist がより小さい隣接ノード）という漸化式で正しく数え上げることができます。計算量は O(E log V) です。

実装例

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

from collections import defaultdict
from heapq import heappop, heappush


def solve(n, edges):
    graph = defaultdict(dict)
    for u, v, w in edges:
        graph[u][v] = w
        graph[v][u] = w

    paths = [0] * (n + 1)
    paths[n] = 1
    dists = [-1] * (n + 1)
    q = [(0, n)]

    while q:
        dist, node = heappop(q)
        if dists[node] != -1:
            continue

        dists[node] = dist
        for v, w in graph[node].items():
            if dists[v] == -1:
                heappush(q, (dist + w, v))
            elif dists[v] < dists[node]:
                paths[node] += paths[v]

        if node == 1:
            return paths[node] % (10**9 + 7)

    return 0


n = 5
edges = [(1,2,3),(1,3,3),(2,3,1),(1,4,2),(5,2,2),(3,5,1),(5,4,10)]
print(solve(n, edges))

入力

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

出力

3