🌲

優先度付きキュー (ヒープ)

2023/08/16に公開

Ruby

コード

コード (競プロコピペ用)

heap.rb

require "forwardable"

class Heap
  ROOT = 0
  private_constant :ROOT

  class << self
    def [](*ary, **options, &block)
      heapify!(ary, **options, &block)
    end

    def heapify!(...)
      new(...).heapify!
    end

    def heapify(ary, ...)
      heapify!(ary.dup, ...)
    end
  end

  extend Forwardable
  def_delegators :@values, :size, :empty?, :clear, :to_a, :to_ary

  attr_accessor :values

  def initialize(values = [], reverse: false, &compare_by)
    @values = values
    @reverse = reverse
    @compare_by = compare_by || -> e { e }
  end

  def push(*value)
    value.each do |value|
      @values << value
      shift_up
    end
  end

  def peek
    @values.first
  end

  def pop(size = nil)
    if @values.empty?
      return
    end

    if size
      return [size, @values.size].min.times.collect { pop }
    end

    swap(ROOT, -1)
    @values.pop.tap do
      shift_down(ROOT)
    end
  end

  def pop_all
    pop(Float::INFINITY)
  end

  def heapify!
    last_parent.downto(ROOT) { |i| shift_down(i) }
    self
  end

  private

  def shift_up(child = @values.size - 1)
    parent = child_to_parent(child)
    if parent < ROOT
      return
    end
    if compare_lt(parent, child)
      return
    end
    swap(parent, child)
    shift_up(parent)
  end

  def shift_down(parent, size: @values.size)
    l = parent_to_left(parent)
    r = parent_to_right(parent)
    i = parent
    if l < size && compare_lt(l, i)
      i = l
    end
    if r < size && compare_lt(r, i)
      i = r
    end
    if i != parent
      swap(parent, i)
      shift_down(i, size: size)
    end
  end

  def swap(i, j)
    @values[i], @values[j] = @values[j], @values[i]
  end

  def compare_lt(parent, child)
    a = @compare_by[@values[parent]]
    b = @compare_by[@values[child]]
    v = a <=> b
    if v == 0
      return false
    end
    (v < 0) ^ @reverse
  end

  def last_parent
    @values.size / 2 - 1
  end

  def parent_to_left(parent)
    parent * 2 + 1
  end

  def parent_to_right(parent)
    parent * 2 + 2
  end

  def child_to_parent(child)
    (child - 1) / 2
  end
end

class Heap
  class << self
    def attach_to(ary, **options, &block)
      new(ary, **options, &block)
    end
  end

  def <<(...)
    push(...)
  end

  def add(...)
    push(...)
  end

  def unshift(...)
    push(...)
  end

  def top
    peek
  end

  def root
    peek
  end

  def poll(...)
    pop(...)
  end

  def shift(...)
    pop(...)
  end

  def reverse!
    @reverse = !@reverse
    heapify!
  end

  def reverse
    self.class.heapify!(@values.dup, reverse: !@reverse, &@compare_by)
  end

  def dup
    self.class.heapify!(@values.dup, reverse: @reverse, &@compare_by)
  end

  def top_view(parent = ROOT)
    value = @values[parent]
    unless value
      return
    end
    [
      top_view(parent_to_left(parent)),
      value,
      top_view(parent_to_right(parent)),
    ].compact
  end

  def inspect
    o = []
    if empty?
      o << "(empty)"
    else
      o << top_view.inspect
    end
    if @reverse
      o << "(reverse)"
    end
    o * " "
  end
end

class Heap
  class << self
    def sort!(ary, reverse: false, &block)
      ary.tap do
        heap = heapify!(ary, reverse: !reverse, &block)
        heap.pop_all_and_insert_to_self
      end
    end

    def sort(ary, ...)
      sort!(ary.dup, ...)
    end
  end

  def pop_all_and_insert_to_self
    (@values.size - 1).downto(1) do |i|
      swap(ROOT, i)
      shift_down(ROOT, size: i)
    end
  end
end

PriorityQueue = Heap

if $0 == __FILE__ && ENV["ATCODER"] != "1"
  require "rspec/autorun"
  RSpec.configure do |config|
    config.expect_with :test_unit
  end

  describe Heap do
    it "works" do
      assert { Heap[3, 1, 2].pop_all == [1, 2, 3] }
      assert { Heap[3, 1, 2, reverse: true].pop_all == [3, 2, 1] }
    end

    it "block" do
      assert { Heap[3, 1, 2] { |e| -e }.pop_all == [3, 2, 1] }
    end

    it "push / add / << / unshift" do
      h = Heap[]
      h << 3
      h.add(1)
      h.push(2)
      h.unshift(4)
      assert { h.pop_all == [1, 2, 3, 4] }
    end

    it "top / peek / root" do
      assert { Heap[3, 1, 2].top == 1 }
      assert { Heap[3, 1, 2].peek == 1 }
      assert { Heap[3, 1, 2].root == 1 }
    end

    it "pop / poll / shift" do
      h = Heap[3, 1, 2]
      assert { h.pop == 1   }
      assert { h.poll == 2  }
      assert { h.shift == 3   }
      assert { h.pop == nil }
    end

    it "扱うオブジェクトは何でもよい" do
      assert { Heap[[3, 0], [1, 0], [2, 0]].pop_all == [[1, 0], [2, 0], [3, 0]] }
      assert { Heap[[3, 0], [1, 0], [2, 0]].reverse.pop_all == [[3, 0], [2, 0], [1, 0]] }
    end

    it "shift_down が正しく動作する" do
      values = [1, 5, 0, 2, 2, 6, 4, 7]
      assert { Heap[*values].pop_all == values.sort }
    end

    it "dup してもブロックを使って比較する" do
      heap = Heap.new { { 1 => 2, 2 => 1 } }
      heap.push(1, 2)
      assert { heap.dup.pop_all == [2, 1] }
      assert { heap.pop_all == [2, 1] }
      assert { heap.empty? }
    end

    it "top_view" do
      assert { Heap[3, 1, 2].top_view == [[3], 1, [2]] }
    end

    it "inspect" do
      assert { Heap[3, 1, 2].inspect == "[[3], 1, [2]]" }
      assert { Heap[].inspect        == "(empty)"       }
    end
  end
end
# >> ..........
# >> 
# >> Finished in 0.01259 seconds (files took 0.07097 seconds to load)
# >> 10 examples, 0 failures
# >>

仕様

Heap と PriorityQueue は同じ

PriorityQueue == Heap  # => true

優先度付きキューの実装にはいくつかあり、その一つがヒープなので優先度付きキューの意味で Heap クラスが剥き出しになると、DDD 的には避けた方がよいとされている技術駆動命名になってしまってよろしくない。そういうときは PriorityQueue の名前で使いたい。そのためにエイリアスを用意してある。

ただこの記事では字面的にシンプルなので Heap を用いる。

値を入れて取り出す

h = Heap[]
h.push(3)
h.push(1)
h.push(2)
h.pop  # => 1
h.pop  # => 2
h.pop  # => 3

poll, shift のエイリアスがある

h = Heap[]
h.push(3)
h.push(1)
h.push(2)
h.poll   # => 1
h.poll   # => 2
h.shift  # => 3

Java の標準クラスライブラリ java.util.PriorityQueue を参考にして入れたけど、頂上から取り出すイメージとは遠い気がしている。

まとめて作る

Heap[3, 1, 2]  # => [[3], 1, [2]]

push 時にまとめて入れる

h = Heap[]
h.push(3, 1, 2)

<<, add, unshift のエイリアスがある

h = Heap[]
h << 3
h.add(1)
h.push(2)
h.unshift(4)
h.pop_all  # => [1, 2, 3, 4]

何個か取り出す

Heap[3, 1, 2].pop(2)  # => [1, 2]

全部取り出す

Heap[3, 1, 2].pop_all  # => [1, 2, 3]

sort_by 風なブロックを使う

Heap[3, 1, 2] { |e|  e }.pop_all  # => [1, 2, 3]
Heap[3, 1, 2] { |e| -e }.pop_all  # => [3, 2, 1]

単に符号を逆にした場合 reverse: true を指定したのと同じ結果になる。

配列を入れても動く

heap = Heap[]
heap << [3, 0]
heap << [1, 0]
heap << [2, 0]
heap.pop_all  # => [[1, 0], [2, 0], [3, 0]]

<=> メソッドに反応するなら何でもよい。

new はあまり使ってほしくない

Heap.new(ary) の意味がコードを見てもよくわからないので当初、禁止にしていたが、Heap[] { |e| } と書くのがちょっと気持ち悪かったので使えるようにした。

Heap.new                 # => (empty)
Heap.new { |e| -e }      # => (empty)
Heap.new(reverse: true)  # => (empty) (reverse)

ヒープ化済みの配列を渡すときはなるべく意味が伝わるように attach_to を使う。from_xxx みたいな名前の方がいいかもしれないが思いつかない。

配列をヒープ化する

a = [3, 1, 2]
h = Heap.heapify(a)  # => [[3], 1, [2]]
a == [3, 1, 2]       # => true

元は壊れない。

配列をヒープ化する (破壊的)

a = [3, 1, 2]
h = Heap.heapify!(a)  # => [[3], 1, [2]]
a != [3, 1, 2]        # => true

渡した配列を in-place でヒープ化する。Ruby の慣習としてクラスメソッドに ! がついているものを見たことがないが、とにかく危険、引数をぶっ壊すので注意せよの意味合いだけでつけてみた。機能的には Python の heapq.heapify(ary) と同じ。

ヒープ済みの配列を渡す (参照渡し)

a = [3, 1, 2]
Heap.heapify!(a)

h = Heap.attach_to(a)
h  # => [[3], 1, [2]]

heapify! でヒープ化してある並びの配列を受けとるだけ。ただヒープ化されているかどうかを使う側が把握しておかないといけないのが負担になるからいれたくなかった。でもその点で言えば Array#bsearch も同じなので仕方なく入れてみた。

空か？

Heap[].empty?  # => true

要素数

Heap[3, 1, 2].size  # => 3

消す

h = Heap[3, 1, 2]
h.clear
h.empty?  # => true

頂上を参照する

Heap[3, 1, 2].peek  # => 1
Heap[3, 1, 2].top   # => 1
Heap[3, 1, 2].root  # => 1

値がなくても騷がない

h = Heap[]
h.peek  # => nil
h.pop   # => nil

peek! や pop! もあった方がいいのかもしれないけどそんなに困ってないので入れてない。

木を真上から見る

主にデバッグ用

h = Heap[3, 1, 2]
h.top_view  # => [[3], 1, [2]]

親1の左右に 3 と 2 の子がいるのがわかる。

h.push(6, 4, 5, 7)
h.top_view  # => [[[6], 3, [4]], 1, [[5], 2, [7]]]

さらにぶら下がって 3 と 2 にも二人の子がいるのがわかる。

内部配列を参照する

Heap[3, 1, 2].values  # => [1, 3, 2]
Heap[3, 1, 2].to_a    # => [1, 3, 2]

Heap.attach_to を使ってセットしていない限り、ここで参照する値は必ずヒープ化している。

inspect する

Heap[].inspect                   # => "(empty)"
Heap.new(reverse: true).inspect  # => "(empty) (reverse)"
Heap[3, 1, 2].inspect            # => "[[3], 1, [2]]"

反転する

h = Heap[3, 1, 2]  # => [[3], 1, [2]]
x = h.reverse      # => [[1], 3, [2]] (reverse)
h                  # => [[3], 1, [2]]
x.pop_all          # => [3, 2, 1]

内部の配列が単に reverse するのではなく昇順・降順を入れ替えて構築し直す。

反転する (破壊的)

h = Heap[3, 1, 2]  # => [[3], 1, [2]]
h.reverse!         # => [[1], 3, [2]] (reverse)
h.pop_all          # => [3, 2, 1]

in-place で行う。

複製する

h = Heap[3, 1, 2]           # => [[3], 1, [2]]
x = h.dup                   # => [[3], 1, [2]]
h.object_id != x.object_id  # => true

最初から降順で作る

Heap[3, 1, 2, reverse: true].pop_all  # => [3, 2, 1]

Heap.heapify([3, 1, 2], reverse: true).pop_all  # => [3, 2, 1]

Heap.heapify!([3, 1, 2], reverse: true).pop_all  # => [3, 2, 1]

ソート

昇順

a = [3, 1, 2]   # => [3, 1, 2]
Heap.sort(a)    # => [1, 2, 3]
a == [3, 1, 2]  # => true

降順

a = [3, 1, 2]                # => [3, 1, 2]
Heap.sort(a, reverse: true)  # => [3, 2, 1]
a == [3, 1, 2]               # => true

昇順 (破壊的)

a = [3, 1, 2]   # => [3, 1, 2]
Heap.sort!(a)   # => [1, 2, 3]
a == [1, 2, 3]  # => true

降順 (破壊的)

a = [3, 1, 2]                 # => [3, 1, 2]
Heap.sort!(a, reverse: true)  # => [3, 2, 1]
a == [3, 2, 1]                # => true

これは動作テスト的な意味で入れているだけで一度きりのソートをするなら Array#sort の方がはるかに速い。

問題を解く

競技プログラミングの鉄則演習問題集 A53 - Priority Queue

優先度付きキュー自体の push, top, pop が正しく動作しているか確認する問題。

N = gets.to_i                                       # => 3
Q = N.times.collect { gets.split.collect(&:to_i) }  # => [[1, 2420], [1, 1650], [2]]
queue = PriorityQueue.new
Q.each do |command, gold|
  case command
  when 1
    queue << gold
  when 2
    p queue.top
  when 3
    queue.pop
  end
end