Ruby と Python で解く AtCoder CODE FESTIVAL 2016 qual C B 優先度付きキュー

はじめに

AtCoder Problems の Recommendation を利用して、過去の問題を解いています。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

AtCoder CODE FESTIVAL 2016 qual C B - K個のケーキ
Difficulty: 905

今回のテーマ、優先度付きキュー

以前解きました Ruby と Python と Java で解く AtCoder ABC141 D 優先度付きキュー の応用版といった問題です。

与えられた配列の最大値と2番目の最大値をキューから取り出し、1を減算してキューに戻します。これを繰返しキューの中身が1つ以下になったところで、ループを抜け出します。

尚、数学的解法もあります。

Ruby 数学的解法

ruby.rb

k = gets.to_i
a = gets.split.map(&:to_i).max
puts [0, 2 * a - k - 1].max

最大値に注目して解きますが、スッキリ感は凄いですね。

Ruby Heap

ruby.rb

class Heap
  attr_reader :size
  def initialize(up: false)
    @up = up
    @heap = []
    @size = 0
  end
  def sum
    x = 0
    @size.times do |i|
      x += @heap[i]
    end
    x
  end
  def push(n)
    n = -n if @up
    i = @size
    while i > 0
      pid = (i - 1) / 2
      break if n >= @heap[pid]
      @heap[i] = @heap[pid]
      i = pid
    end
    @heap[i] = n
    @size += 1
  end
  def pop
    return nil if @size == 0
    top = @heap[0]
    @size -= 1
    n = @heap[@size]
    i = 0
    while i * 2 + 1 < @size
      cid1 = i * 2 + 1
      cid2 = cid1 + 1
      if cid2 < @size && @heap[cid2] < @heap[cid1]
        cid1 = cid2
      end
      break if @heap[cid1] >= n
      @heap[i] = @heap[cid1]
      i = cid1
    end
    @heap[i] = n
    if @up
      -top
    else
      top
    end
  end
end

gets
a = gets.split.map(&:to_i)
h = Heap.new(up: true)
a.each do |x|
  h.push(x)
end
while h.size > 1
  u = h.pop
  v = h.pop
  h.push(u - 1) if u - 1 > 0
  h.push(v - 1) if v - 1 > 0
end
if h.size == 0
  puts "0"
else
  puts h.pop - 1
end

up.rb

  def initialize(up: false)
    @up = up

前回 のクラスを改良して、最大値と最小値をフラグで切り替えられるようにしました。

heap.rb

while h.size > 1
  u = h.pop
  v = h.pop
  h.push(u - 1) if u - 1 > 0
  h.push(v - 1) if v - 1 > 0
end

2つ取り出して減算した後、1以上であればキューに戻します。

Python　数学的解法

python.py

from sys import stdin

def main():
    input = stdin.readline
    k, _ = map(int, input().split())
    a = max(map(int, input().split()))
    print(max(0, 2 * a - k - 1))
main()

Python heapq

python.py

from sys import stdin

def main():
    import heapq
    input = stdin.readline
    input()
    a = [-1 * int(i) for i in input().split()]
    heapq.heapify(a)
    while len(a) > 1:
        u = heapq.heappop(a)
        v = heapq.heappop(a)
        if u + 1 < 0:
            heapq.heappush(a, u + 1)
        if v + 1 < 0:
            heapq.heappush(a, v + 1)
    if len(a) == 0:
        print(0)
    else:
        print(abs(a[0] + 1))
main()

演算が複雑な場合、符号変換用の関数を準備した方がいいと思われます。

	Ruby 数学的解法	Ruby Heap	Python 数学的解法	Python heapq
コード長 (Byte)	76	1120	184	458
実行時間 (ms)	7	26	17	22
メモリ (KB)	1788	1788	2940	3064

まとめ

• AtCoder CODE FESTIVAL 2016 qual C B を解いた
• Ruby に詳しくなった
• Python に詳しくなった

参照したサイト
[競プロ用]優先度付きキュー(heapq)まとめ