Rubyで優先度つきキューを実装

競技プログラミングをやっていて、きちんと準備しておかないとキツいと感じたので。

Rubyでの優先度つきキューは、検索すればたくさんの人の実装が見つかるので自作する必要は無いのだが、練習も兼ねて自分ならどう実装するか考えてみることにした。

## 欲しい機能
要素の追加 `#<<` （または `#push` ）と先頭要素の削除 `#shift` が O(log N) で実行できること。

あとは O(1) で実行できる簡単で非破壊的な処理 `#first`, `#size`, `#empty?` も念のため入れておく。（メソッド名は `Array` に合わせたもの）

---

加えて、 `queue << [priority, value]` のように **`Array` のデータを入れる**ことがあるため、それに対応したい。

- キューに `Array` を入れても動作すること
	- `#<=>` はあるが `Comparable` ではない[^array-comp]ため、多くの比較演算子が使えない
- 昇順・降順など、優先度の定義を好きに決められること
	- 些細な違いのために一々コードを書き換えたくない
	- `#sort_by` のようにソートキーを決めたり、 `#sort` のように2値の比較方法自体を指定したりしたい
		- 要素が `[priority, value]` の場合、 `#first` で評価すればいいのではという思いがある

【追記】
この目的で `Array` を使ったり `#<=>` で比較するのは実行時間がのびる（TLEになりやすい）ことがわかったので、 `queue.push(priority, value)` とデータを分離した形式も作ることにした。

[^array-comp]: `Array` は*常に*比較可能というわけではないので `Comparable` にしていなく、 `#<=>` の追加はソートのための妥協点らしい。 https://bugs.ruby-lang.org/issues/5574

## 実装
優先度の決め方は、 `.new` にブロックで渡すことにした。

- 無指定なら `#<=>` を使い、 `Array` を入れる場合も動作するようにする
- ブロックが2引数なら、比較方法と判断してそのまま使う
- ブロックが1引数なら、ソートキーを求めて比較するprocを作る

---

単純な「[二分ヒープ](https://ja.wikipedia.org/wiki/二分ヒープ)」で実装する。配列のインデックスで二分木の親子関係を表現できるので、データ構造が簡潔になる。

要素の追加および先頭要素の削除の際に、親子のノード間で順序を正す作業が発生する。それをプライベートメソッド `#update` に纏めたものが以下のコード。（各操作で別々に実装すればもっと最適化できる）

```ruby:priority_queue.rb
class PriorityQueue
	def initialize(&compare)
		# ヒープ構造を表す配列（インデックス0は不使用）
		@heap = [nil]

		# 要素の比較方法
		compare ||= ->(x, y) { x <=> y }
		@comp = case compare.arity
		when 1 then ->(x, y) { compare[x] <=> compare[y] }
		when 2 then compare
		else        raise ArgumentError
		end
	end

	def <<(obj)
		# 一旦末尾に加える
		@heap << obj

		# ヒープの順序を正す（葉から根へ）
		k = @heap.size - 1
		while (k >>= 1) > 0 && update(k); end

		# queue << obj1 << obj2 と連続できるようselfを返す
		return self
	end

	def shift
		return nil if @heap.size == 1
		return @heap.pop if @heap.size == 2

		# 先頭を抜き取り、末尾を一旦先頭に置く
		result = @heap[1]
		@heap[1] = @heap.pop

		# ヒープの順序を正す（根から葉へ）
		k = 1
		while (k = update(k)); end

		return result
	end

	def first
		return @heap[1]  # 要素が無ければnil
	end

	def size
		return @heap.size - 1
	end

	def empty?
		return @heap.size == 1
	end

	private

	# 位置kとその子ノードの間で順序を正す
	# 入れ替えたら子のインデックスを、入れ替えなかったらnilを返す
	def update(k)
		l, r = k << 1, k << 1 | 1
		up, left, right = @heap.values_at(k, l, r)
		return nil unless left  # 子が無い

		# 左右の子ノードのうち値が小さいほうが交換候補
		if right && @comp[left, right] > 0
			return nil if @comp[up, right] <= 0
			@heap[k], @heap[r] = right, up
			return r
		else
			return nil if @comp[up, left] <= 0
			@heap[k], @heap[l] = left, up
			return l
		end
	end
end

if $0 == __FILE__
	# ヒープソート（降順）でテスト
	data = Array.new(200_000) { rand(1_000_000) }
	pq = PriorityQueue.new { |x| -x }
	data.each { |x| pq << x }
	sorted = data.map { pq.shift }

	# p data, sorted, pq.shift
	raise "Wrong sort!" if data.sort_by { |x| -x } != sorted
	raise "Queue is not empty!" if pq.shift
	puts "OK."
end
```

AtCoderのコードテストで実行して、20万件の整数のソートが 1900ms 台で完了した。要素の追加・削除ともきちんと O(log N) になっていると思われる。

また、「降順」の指定を `{ |x, y| y <=> x }` にしたら 1200ms 台まで縮んだ。ソートキーを指定する方法は、この実装のままだと実用的でないかもしれない。

## 実装（キー分離版）
`queue.push(priority, value)` と使うことで、 `Array` の作成や `#<=>` の使用を回避したバージョン。新しく `@objs` を用意し、キーに対応するオブジェクトを `@heap` と同じインデックスに管理するようにした。

<details><summary>コードを表示</summary><div>

```ruby:priority_queue.rb
class PriorityQueue
	def initialize
		# ヒープ構造を表す配列（インデックス0は不使用）
		@heap = [nil]
		@objs = [nil]
	end

	def push(key, obj)
		# 一旦末尾に加える
		@heap << key
		@objs << obj

		# ヒープの順序を正す（葉から根へ）
		k = @heap.size - 1
		while (k >>= 1) > 0 && update(k); end

		# メソッドを連続できるようselfを返す
		return self
	end

	def shift
		return nil if @heap.size == 1
		return @heap.pop, @objs.pop if @heap.size == 2

		# 先頭を抜き取り、末尾を一旦先頭に置く
		key0, obj0 = @heap[1], @objs[1]
		@heap[1], @objs[1] = @heap.pop, @objs.pop

		# ヒープの順序を正す（根から葉へ）
		k = 1
		while (k = update(k)); end

		return key0, obj0
	end

	def first
		return nil if @heap.size == 1
		return @heap[1], @objs[1]
	end

	def size
		return @heap.size - 1
	end

	def empty?
		return @heap.size == 1
	end

	private

	# 位置kとその子ノードの間で順序を正す
	# 入れ替えたら子のインデックスを、入れ替えなかったらnilを返す
	def update(k)
		l, r = k << 1, k << 1 | 1
		up, left, right = @heap.values_at(k, l, r)
		return nil unless left  # 子が無い

		# 左右の子ノードのうち値が小さいほうが交換候補
		if right && left > right
			return nil if up <= right
			@heap[k], @heap[r] = right, up
			@objs[k], @objs[r] = @objs[r], @objs[k]
			return r
		else
			return nil if up <= left
			@heap[k], @heap[l] = left, up
			@objs[k], @objs[l] = @objs[l], @objs[k]
			return l
		end
	end
end

if $0 == __FILE__
	# ヒープソート（降順）でテスト
	data = Array.new(200_000) { rand(1_000_000) }
	pq = PriorityQueue.new
	data.each { |x| pq.push(-x, x) }
	sorted = data.map { pq.shift[1] }

	# p data, sorted, pq.shift
	raise "Wrong sort!" if data.sort_by { |x| -x } != sorted
	raise "Queue is not empty!" if pq.shift
	puts "OK."
end
```

</div></details>

<details><summary>差分のみ表示</summary><div>

```diff_ruby:差分
@@ -1,20 +1,14 @@
 class PriorityQueue
-	def initialize(&compare)
+	def initialize
 		# ヒープ構造を表す配列（インデックス0は不使用）
 		@heap = [nil]
+		@objs = [nil]
-
-		# 要素の比較方法
-		compare ||= ->(x, y) { x <=> y }
-		@comp = case compare.arity
-		when 1 then ->(x, y) { compare[x] <=> compare[y] }
-		when 2 then compare
-		else        raise ArgumentError
-		end
 	end
 
-	def <<(obj)
+	def push(key, obj)
 		# 一旦末尾に加える
-		@heap << obj
+		@heap << key
+		@objs << obj
 
 		# ヒープの順序を正す（葉から根へ）
 		k = @heap.size - 1
@@ -26,21 +20,22 @@
 
 	def shift
 		return nil if @heap.size == 1
-		return @heap.pop if @heap.size == 2
+		return @heap.pop, @objs.pop if @heap.size == 2
 
 		# 先頭を抜き取り、末尾を一旦先頭に置く
-		result = @heap[1]
+		key0, obj0 = @heap[1], @objs[1]
-		@heap[1] = @heap.pop
+		@heap[1], @objs[1] = @heap.pop, @objs.pop
 
 		# ヒープの順序を正す（根から葉へ）
 		k = 1
 		while (k = update(k)); end
 
-		return result
+		return key0, obj0
 	end
 
 	def first
+		return nil if @heap.size == 1
-		return @heap[1]  # 要素が無ければnil
+		return @heap[1], @objs[1]
 	end
 
 	def size
@@ -61,13 +56,15 @@
 		return nil unless left  # 子が無い
 
 		# 左右の子ノードのうち値が小さいほうが交換候補
-		if right && @comp[left, right] > 0
+		if right && left > right
-			return nil if @comp[up, right] <= 0
+			return nil if up <= right
 			@heap[k], @heap[r] = right, up
+			@objs[k], @objs[r] = @objs[r], @objs[k]
 			return r
 		else
-			return nil if @comp[up, left] <= 0
+			return nil if up <= left
 			@heap[k], @heap[l] = left, up
+			@objs[k], @objs[l] = @objs[l], @objs[k]
 			return l
 		end
 	end
@@ -76,9 +73,9 @@
 if $0 == __FILE__
 	# ヒープソート（降順）でテスト
 	data = Array.new(200_000) { rand(1_000_000) }
-	pq = PriorityQueue.new { |x| -x }
+	pq = PriorityQueue.new
-	data.each { |x| pq << x }
+	data.each { |x| pq.push(-x, x) }
-	sorted = data.map { pq.shift }
+	sorted = data.map { pq.shift[1] }
 
 	# p data, sorted, pq.shift
 	raise "Wrong sort!" if data.sort_by { |x| -x } != sorted
```

</div></details>

テストのソートにかかった時間は 1100ms 台になった。 `@obj` の管理が増えているものの、キーの比較が簡潔になった利点のほうが上回るらしい。

## 参考

- [二分ヒープ - Wikipedia](https://ja.wikipedia.org/wiki/二分ヒープ)
- [class Array (Ruby リファレンスマニュアル)](https://docs.ruby-lang.org/ja/latest/class/Array.html)

実際のところ、大学の演習でヒープソートを学んだ経験が一番活きている。