グラフアルゴリズムの理解のためにBellman–Ford法をRubyで実装してみる

優先度付きキューにも使われる二分ヒープ構造をRubyで実装してみる - $shibayu36->blog;に引き続き。アルゴリズム図鑑を眺めていて、グラフ系のアルゴリズムを全く知らないことに気づいた。そこでこの本に載っていたBellman-Ford法についてRubyで実装してみたのでメモ。

重み付きグラフ

Bellman-Ford法の前に重み付きのグラフの構造を作っておく必要がある。以下のように実装をした。無向グラフも有向グラフも扱えるように作った。

class WeightedGraph
  def initialize
    @graph = Hash.new { |h, k| h[k] = [] }
  end

  def add_edge(from, to, weight)
    @graph[from].push(Edge.new(from, to, weight))
    @graph[to].push(Edge.new(to, from, weight))
  end

  def add_directed_edge(from, to, weight)
    @graph[from].push(Edge.new(from, to, weight))
  end

  def points
    @graph.keys
  end

  def edges
    @graph.values.flat_map do |edges|
      edges.map do |edge|
        { from: edge.from, to: edge.to, weight: edge.weight }
      end
    end
  end

  class Edge
    attr_reader :from, :to, :weight

    def initialize(from, to, weight)
      @from = from
      @to = to
      @weight = weight
    end
  end
end

Bellman-Fordアルゴリズム

Bellman-Ford法はグラフの最短経路問題を解くアルゴリズムだ。ダイクストラ法の方が高速だが、Bellman-Ford法は負の経路がある場合にも使える。

実装については以下を参考にした。 * ベルマン–フォード法 - Wikipedia * Bellman-Ford法の解説 - Qiita * 【アルゴリズム】ベルマンフォード法(BellmanFord)

実装

class BellmanFord
  def initialize(graph, from, to)
    @graph = graph
    @from = from
    @to = to
  end

  def shortest_cost
    @last_updated = Hash.new(nil)

    @cost = Hash.new(Float::INFINITY)
    @cost[@from] = 0

    (0..@graph.points.length - 1).each do |i|
      is_changed = false

      @graph.edges.each do |edge|
        candidate_cost = @cost[edge[:from]] + edge[:weight]
        if candidate_cost < @cost[edge[:to]]
          @cost[edge[:to]] = candidate_cost
          @last_updated[edge[:to]] = edge[:from]
          is_changed = true

          # negative loop detection.
          # This algorithm must finish within n - 1 iterations.
          # If the N-th iteration change cost, there must be a negative loop.
          if i == @graph.points.length - 1
            raise 'negative loop detected'
          end
        end
      end

      break unless is_changed
    end

    @cost[@to]
  end

  def shortest_path
    shortest_cost
    path = [@to]
    loop do
      break if @last_updated[path[0]].nil?

      path.unshift(@last_updated[path[0]])
    end
    path
  end
end

テスト

describe WeightedGraph do
  it 'calculates shortest cost' do
    graph = WeightedGraph.new
    graph.add_edge('A', 'B', 9)
    graph.add_edge('A', 'C', 2)
    graph.add_edge('B', 'C', 6)
    graph.add_edge('B', 'D', 3)
    graph.add_edge('B', 'E', 1)
    graph.add_edge('C', 'D', 2)
    graph.add_edge('C', 'F', 9)
    graph.add_edge('D', 'E', 5)
    graph.add_edge('D', 'F', 6)
    graph.add_edge('E', 'F', 3)
    graph.add_edge('E', 'G', 7)
    graph.add_edge('F', 'G', 4)
    expect(BellmanFord.new(graph, 'A', 'G').shortest_cost).to eq(14)
    expect(BellmanFord.new(graph, 'A', 'G').shortest_path).to eq(%w[A C D F G])
  end

  it 'also calculates shortest cost when graph is directed' do
    graph = WeightedGraph.new
    graph.add_directed_edge(0, 2, 1)
    graph.add_directed_edge(0, 3, 4)
    graph.add_directed_edge(0, 4, 5)
    graph.add_directed_edge(2, 1, 1)
    graph.add_directed_edge(3, 6, 4)
    graph.add_directed_edge(4, 5, 2)
    graph.add_directed_edge(4, 6, 2)
    graph.add_directed_edge(1, 5, 4)
    graph.add_directed_edge(1, 7, 8)
    graph.add_directed_edge(5, 7, 2)
    graph.add_directed_edge(6, 7, 5)
    expect(BellmanFord.new(graph, 0, 7).shortest_cost).to eq(8)
    expect(BellmanFord.new(graph, 0, 7).shortest_path).to eq([0, 2, 1, 5, 7])
  end

  it 'can detect negative loop' do
    graph = WeightedGraph.new
    graph.add_edge('A', 'B', 1)
    graph.add_edge('B', 'C', 3)
    graph.add_edge('C', 'D', -5)
    graph.add_edge('D', 'B', 1)
    graph.add_edge('C', 'E', 2)
    expect { BellmanFord.new(graph, 'A', 'E').shortest_cost }.to raise_error('negative loop detected')
  end
end

まとめ

今回は自分の中でグラフのアルゴリズムの理解を深めるため、Rubyで実装をしてみた。計算を繰り返すと最短経路を求められるというのは面白いと感じた。