欲しい機能

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

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

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

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

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

実装

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

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

単純な「二分ヒープ」で実装する。配列のインデックスで二分木の親子関係を表現できるので、データ構造が簡潔になる。

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

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 と同じインデックスに管理するようにした。

コードを表示

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

差分のみ表示

差分

@@ -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

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

参考

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

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