Go to Qiita Advent Calendar 2024 Top
LoginSignup
2
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

More than 3 years have passed since last update.

@HMMNRST

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

Last updated at Posted at 2021-08-07

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

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

欲しい機能

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

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

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

  • キューに Array を入れても動作すること
    • #<=> はあるが Comparable ではない1ため、多くの比較演算子が使えない
  • 昇順・降順など、優先度の定義を好きに決められること
    • 些細な違いのために一々コードを書き換えたくない
    • #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 の管理が増えているものの、キーの比較が簡潔になった利点のほうが上回るらしい。

参考

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

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

2
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
2
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?