RubyでPriority Queue

Feb 14, 2020 02:42 · 1290 words · 3 minute read

今日の夜、会社の競プロイベント#2が開催された。バーチャルコンテスト形式で10人以上が参加してわいわいやった。途中参加だったけどめちゃくちゃ楽しかった。

問題の中で、愚直にやったらTLEするようなものがあった。計算量に気をつけないといけないやつ。

(この問題 => C - Factory

結論から言うと、Priority Queue(優先度付きキュー)があればシュッと解けるんだけど、Rubyの組み込みライブラリにはなさそうだった。

Priority QueueはHeap(木構造が)あれば実装できるけど、それもなさそうな感じだったので書いてみた。

それがこちら。
https://gist.github.com/akht/3749fac291c527aeb043c40b12c4950a

このブログでは展開されないので中身を貼っておくとこんな感じになってる。

class MinHeap
  def initialize(source)
    @arr = []
    source.each do |e|
      push(e)
    end
  end

  def size
    @arr.size
  end

  def empty?
    @arr.empty?
  end

  # ヒープの最小値を返す(削除しない)
  def top
    @arr.first
  end

  # ヒープに値を追加する
  def push(value)
    @arr << value     # 最後尾に追加
    i = @arr.size - 1 # 追加されたノード番号

    # 親子関係を修復していく
    while i > 0
      parent = (i - 1) / 2 # 親のノード番号

      # 親子関係が逆転してなかったら終了
      break if @arr[parent] <= value

      # 親と自分をswap
      @arr[parent], @arr[i] = @arr[i], @arr[parent]

      i = parent
    end

    # 追加したい値は最終的にこの位置に決定する
    @arr[i] = value
  end

  # ヒープから最小値を取り出す
  def pop
    min = top # 最小値 = 返り値

    # ここからヒープを再構築する
    # ひとまず最後尾をルートノードにもっていくる
    tmp_node = @arr.last
    @arr.pop

    # ルートノードから降ろしていく
    i = 0
    while (i * 2 + 1) < size
      # 左側の子が配列のサイズを超えるまで
      # = 左側の子が配列サイズを超えているということは、
      # そのノードには子が存在しないということになる

      # 左右の子ノードの値を比較して小さい方をmin_childとする
      lhs_child = i * 2 + 1 # 左側の子
      rhs_child = i * 2 + 2 # 右側の子
      min_child = lhs_child # とりあえず左側が小さいと仮定しておく(右側の子が存在しないこともあるため)

      if rhs_child < size
        # 右側の子のインデックスが配列のサイズを超えていないとき
        # = 超えているということは、右側の子は存在しないということ
        # => その場合は左側の子が採用される

        if @arr[lhs_child] > @arr[rhs_child]
          min_child = rhs_child
        end
      end

      # 逆転していなければ終了
      break if @arr[min_child] >= tmp_node

      # 自分のノードを子の値にする
      @arr[i] = @arr[min_child]
      i = min_child
    end

    # 選んだノードは最終的にこの位置に決定(空なら最後の要素だったということなのでセット不要)
    @arr[i] = tmp_node unless empty?

    min
  end
end

厳密に言えば、これはヒープの中でも二分ヒープというやつ(たぶん)。
ルートノードが最小値になるように木を構築してるので最小ヒープと呼ぶ。逆にルートノードが最大値になるやつは最大ヒープと呼ぶ。

つまり、最大/最小値が知りたいときは、ルートノードを参照すればいいので\(O(1)\)ですむ。この実装の場合だと配列を使ってるから0番目のインデックスの値をみればいい。すごい。
新しく要素を追加したり(push)、削除する(pop)のは\(O(logN)\)で行える。
これはノードが増えたり減ったりするとヒープを再構築する必要があって、それには二分木の高さ = \(logN\)(要素数N)のステップが必要なため。

というわけで、これがあると優先度付きキューもシュッと実装できて(というかそのまま使える)、問題もあっさり解けてしまう。 なんて便利な世の中になったんだろう。
https://atcoder.jp/contests/code-thanks-festival-2017-open/submissions/10080104