Zenn
🌳

プリム法

2024/02/16に公開
Ruby
AtCoder
競プロ
最小全域木
プリム法
tech

とは？

最小全域木を作るアルゴリズムの一つでクラスカル法に比べて人気がない。

アルゴリズム

  1. どこでもいいのでスタート地点をコスト0として優先度付きキューに入れる
  2. キューから最小コストで行ける場所を取り出す
  3. 初めての場所であればコストを確定し、そこから初めて進める次の場所をキューに入れる
  • 2〜3 を繰り返す
  • 確定したコストの合計が最小コストになる

具体例

次の路線図は無駄があるのでプリム法で適切な路線図と最小コストを求める。

辺 コスト
A-B 3
B-C 2
C-A 1

仮に A からスタートすると、

取り出す コスト確定 次の目的地とコスト
A 0 0
→ B 3
→ C 1
C 1 1
→ B 2
→ A 1 (巡回済み)
B 2 2
→ A 3 (巡回済み)
→ C 2 (巡回済み)
B 3 (巡回済み)

となって 0 + 1 + 2 = 3 が答えになる。

実装例

edges = [["A", "B", 3], ["B", "C", 2], ["C", "A", 1]]

G = Hash.new { |h, k| h[k] = {} }
edges.each do |x, y, c|
  G[x][y] = c
  G[y][x] = c
end
G                     # => {"A"=>{"B"=>3, "C"=>1}, "B"=>{"A"=>3, "C"=>2}, "C"=>{"B"=>2, "A"=>1}}

start = G.keys.first  # => "A"

visited = {}
total = 0

heap = Heap.new
heap << [0, start]

until heap.empty?
  cost, from = heap.pop
  visited[from] and next
  visited[from] = true
  total += cost
  G[from].each do |to, cost|
    visited[to] and next
    heap << [cost, to]
  end
end

total                 # => 3

最小全域森には適さない

たとえば A-B C-D となっているとき A から始めてしまうと C-D のことはまったくわからないまま終わってしまう。

過去問を解く

典型アルゴリズム問題集 F - 最小全域木問題

https://atcoder.jp/contests/typical-algorithm/tasks/typical_algorithm_f

これは実装例を標準入力に対応しただけで通る。実行時間は 1338 ms で、クラスカル法の 375 ms に比べると異様に遅いが、2 秒の制限にはまだ余裕がある。

コード (AC 1338 ms)
N, M = gets.split.collect(&:to_i)                       # => [3, 3]
edges = M.times.collect { gets.split.collect(&:to_i) }  # => [[2, 3, 7202], [1, 2, 629], [1, 3, 3497]]

G = Hash.new { |h, k| h[k] = {} }
edges.each do |x, y, c|
  G[x][y] = c
  G[y][x] = c
end
G                                                       # => {2=>{3=>7202, 1=>629}, 3=>{2=>7202, 1=>3497}, 1=>{2=>629, 3=>3497}}

start = G.keys.first                                    # => 2

visited = {}
total = 0

heap = Heap.new
heap << [0, start]

until heap.empty?
  cost, from = heap.pop
  visited[from] and next
  visited[from] = true
  total += cost
  G[from].each do |to, cost|
    visited[to] and next
    heap << [cost, to]
  end
end

p total                                                 # => 4126

競技プログラミングの鉄則 演習問題集 A67 - MST

https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_bo

上と同じコードで提出すると 1110 ms で TLE になる。クラスカル法の場合 382 ms で通ったので気づかなかったが、どうやら実行時間が 1 秒に設定されているらしい。そこで競プロ攻略のための小細工が必要になってくる。具体的には、

  • 節を 0 ベースの整数に揃える
  • グラフ G と巡回済み情報 visited の型を配列にする
  • キュー要素の比較方法をブロックで明示する

とすることで 232 ms だけ速くなりなんとか通る。

最適化済みコード (AC 878 ms)
N, M = gets.split.collect(&:to_i)  # => [3, 3]
edges = M.times.collect { gets.split.collect(&:to_i) }

G = Array.new(N) { [] }
edges.each do |x, y, c|
  G[x.pred] << [y.pred, c]
  G[y.pred] << [x.pred, c]
end

start = 0

visited = Array.new(N)
total = 0

heap = Heap.new { |cost, _| cost }
heap << [0, start]

until heap.empty?
  cost, from = heap.pop
  visited[from] and next
  visited[from] = true
  total += cost
  G[from].each do |to, cost|
    visited[to] and next
    heap << [cost, to]
  end
end

p total                            # => 4126

関連

https://zenn.dev/megeton/articles/f399ad3da537e5

Discussion