Rubyでダイクストラ法（ヒープを使う版）

Rubyでダイクストラ法（ヒープを使わない版）の続き。
Graphはベルマン-フォード法のとこで作ったもの、最小ヒープは先日作ったものを使う。

今度は最小ヒープを優先度付きキューとして使って実装した。

コード

require "./graph"
require "./heap"

class ShortestRouteFinder
  def initialize(graph)
    @graph = graph
  end

  # ダイクストラ法
  def search(from, to)
    p "shortest @vertex_weights from #{from}:"

    # 始点から終点までの経路を計算
    calc_routes(from, to)

    p @vertex_weights
    p @routes
  end

  def init_routes(from)
    @vertex_weights = {}
    @routes = {}
    @graph.vertexes.each_key do |k|
      @vertex_weights[k] = Float::INFINITY
      @routes[k] = [from] # fromから各頂点までの最短の経路で通過する頂点を記録
    end

    # 最小コストの頂点を選択するための優先度付きキュー
    @priority_queue = Heap.new

    # from自身のコストは0
    @vertex_weights[from] = 0
  end

  # 頂点に隣接する頂点までのコストを求めることを、終点の最小コストが求まるまで繰り返す
  def calc_routes(starting_vertex, goal_vertex)
    # 全頂点に対してコスト無限で初期化 経路も初期化
    init_routes(starting_vertex)
    @priority_queue.push([@vertex_weights[starting_vertex], starting_vertex])

    Kernel.loop do
      # 未探索の頂点の中から最小のコストを持つ頂点を次の始点とする
      vertex_weight, current_vertex = @priority_queue.pop

      # 終点が最小のコストなら抜ける
      break if current_vertex == goal_vertex

      # 隣接する頂点のコストを更新 current_vertexの頂点からのびる辺をすべて処理
      @graph.vertexes[current_vertex].adjacents.each do |k, adjacent_weight|
        # 現在の頂点のコストと辺のコストを足して隣接する頂点のコストを求める
        # 新しく求めたコストが小さい場合のみ更新
        new_neighbor_weight = vertex_weight + adjacent_weight
        next if @vertex_weights[k] < new_neighbor_weight

        @vertex_weights[k] = new_neighbor_weight
        @routes[k] = @routes[current_vertex] + [k] # 経路を記録
        @priority_queue.push([new_neighbor_weight, k]) # 次の始点の候補として優先度付きキューに入れる
      end
    end
  end
end

graph = Graph.new
%w[A B C D E F G].each { |key| graph.add_vertex(key) }

graph.connect("A", "B", 2)
graph.connect("A", "C", 5)
graph.connect("B", "D", 1)
graph.connect("B", "E", 3)
graph.connect("C", "B", 6)
graph.connect("C", "F", 8)
graph.connect("D", "E", 4)
graph.connect("E", "G", 9)
graph.connect("F", "G", 7)

route_finder = ShortestRouteFinder.new(graph)
route_finder.search("A", "G")
# {"A"=>0, "B"=>2, "C"=>5, "D"=>3, "E"=>5, "F"=>13, "G"=>14}
# {"A"=>["A"], "B"=>["A", "B"], "C"=>["A", "C"], "D"=>["A", "B", "D"], "E"=>["A", "B", "E"], "F"=>["A", "C", "F"], "G"=>["A", "B", "E", "G"]}

優先度付きキューによる変更点

最小のコストを持つ頂点を求めるのに以前は以下のようにselectとsortを駆使していたが

# 未探索の頂点の中から最少のコストを持つ頂点を次の始点とする
next_current_vertex = @vertex_weights.select{|k,v| @unexpored_vertexes.include?(k) }.sort{|v1, v2| v1[1] <=> v2[1]}.first[0]

優先度付きキューを使うことで以下のようになった。シンプルで計算量も少ない。

vertex_weight, current_vertex = @priority_queue.pop

最小コストの頂点を求める計算量はselectとsortでO(n)+O(n log n)だったものが、最小値の取り出しとヒープの再構築でO(1)+O(log n)と激的に小さくなっている。

また最小コストの頂点を見つけるために@unexpored_vertexesで未探索の頂点を保持していたが、探索候補も優先度付きキューに任せられるようになったので、@unexpored_vertexesは不要となった。

@priority_queue.push([new_neighbor_weight, k]) # 次の始点の候補として優先度付きキューに入れる

さらに再帰を使った方が分かりやすいと思ったので再帰を使っていたが、単純にループで分かりやすく書けるようになった。

適切なデータ構造を使うことで実装がシンプルになり計算量も少なくなるという好例だなと思いました。まる。