【Python】バスの乗り継ぎで最終目的地までの最小コストを求めるプログラムの解説

問題の概要

n × 3 の行列が与えられます。各行は [出発地(src)、目的地(dest)、路線ID(id)] の3つのフィールドで構成されており、そのバスが出発地から目的地へ運行していることを表しています。
新しいバスに乗るたびに1単位の料金がかかりますが、同じバスに乗り続けている間は合計で1単位しか支払いません。このとき、場所0から最終目的地(最大の場所番号)まで移動するために必要な最小コストを求めます。目的地に到達できない場合は -1 を返してください。

入力例

上記の入力の場合、出力は 2 になります。場所0でバス0に乗り、場所3で降りて、そこからバス1に乗り換えて場所5へ向かうためです。

解決のアプローチ

この問題は、各地点をノード、バス路線をエッジとみなしたグラフ探索として捉えることができます。乗車コストが最小となる経路を見つける必要があるため、優先度付きキュー(ヒープ)を用いたダイクストラ法が有効です。状態として「現在位置」と「現在乗っているバス」を管理し、バスを乗り換えるたびにコストを1増やしていきます。

具体的な手順は以下の通りです。

• start を 0、target を行列内の最大の場所に設定します
• 隣接リスト adj を作成し、各路線について adj[src] の末尾に (dest, id) を追加します
• ヒープ hp を (0, start, -1)(初期コスト0・開始位置・未乗車)で初期化します
• 訪問履歴を記録する seen を用意します
• hp が空になるまで以下を繰り返します:
  • ヒープから最小コストの要素 (cost, cur_pos, cur_bus) を取り出します
  • cur_pos が target と一致すれば cost を返します
  • cur_bus がすでに seen[cur_pos] に含まれる場合はスキップします
  • seen[cur_pos] に cur_bus を追加します
  • 隣接する各 (nex_pos, nex_bus) について、バスが変わる場合は next_cost に1を加算し、ヒープに push します
• ループが完了しても目的地に届かなければ -1 を返します

計算量は、辺の数を E とすると O(E log E) 程度になり、効率的に最短コストを求められます。

実装例(Python)

from collections import defaultdict
from heapq import heapify, heappop, heappush
class Solution:
   def solve(self, connections):
      start = 0
      target = max(max(y, x) for y, x, _ in connections)

      adj = defaultdict(list)
      for f, t, id in connections:
         adj[f].append((t, id))

      hp = [(0, start, -1)]
      seen = defaultdict(set)

      while hp:
         cost, cur_pos, cur_bus = heappop(hp)
         if cur_pos == target:
            return cost
         if cur_bus in seen[cur_pos]:
            continue
         seen[cur_pos].add(cur_bus)

         for nex_pos, nex_bus in adj[cur_pos]:
            next_cost = cost
            if nex_bus != cur_bus:
               next_cost += 1
            heappush(hp, (next_cost, nex_pos, nex_bus))

      return -1

ob = Solution()
matrix = [
   [0, 1, 0],
   [1, 2, 0],
   [2, 3, 0],
   [3, 5, 1],
   [5, 0, 2]
]
print(ob.solve(matrix))

入力

matrix = [[0, 1, 0],
[1, 2, 0],
[2, 3, 0],
[3, 5, 1],
[5, 0, 2]]

出力

2