Ruby と Python と Java で解く AtCoder ABC141 D 優先度付きキュー

はじめに

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

今回のお題

AtCoder Beginner Contest D - Powerful Discount Tickets
Difficulty: 826

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

Ruby

操作自体はシンプルで、一番値段の高い品物をキューから取り出し、割引券を一枚適用してキューに戻します。その都度、一番値段の高い品物ついて同様の操作を割引券がなくなるまで行います。
但し、次の様に単にソートするだけの実装では、TLEになります。

ruby_tle.rb

n, m = gets.split.map(&:to_i)
a = gets.split.map(&:to_i)
a.sort_by!{|x| -x}
m.times do
  b = a.shift
  b /= 2
  a << b
  a.sort_by!{|x| -x}  
end
puts a.inject(:+)

優先度付きキューは、ソートに比べて少ない計算量で一番値段の高い品物を調べることができます。
Python ですとheapq、Java ですとPriorityQueueになりますが、Ruby には無いので、Ruby で Priority Queue を実装してみたい のコードをお借りして、少々修正して通しました。

ruby.rb

class PriorityQueue
  def initialize(array = [])
    @data = []
    array.each{|a| push(a)}
  end

  def push(element)
    @data.push(element)
    bottom_up
  end

  def pop
    if size == 0
      return nil
    elsif size == 1
      return @data.pop
    else
      min = @data[0]
      @data[0] = @data.pop
      top_down
      return min
    end
  end

  def size
    @data.size
  end

  private

  def swap(i, j)
    @data[i], @data[j] = @data[j], @data[i]
  end

  def parent_idx(target_idx)
    (target_idx - (target_idx.even? ? 2 : 1)) / 2
  end

  def bottom_up
    target_idx = size - 1
    return if target_idx == 0
    parent_idx = parent_idx(target_idx)
    while (@data[parent_idx] > @data[target_idx])
      swap(parent_idx, target_idx)
      target_idx = parent_idx
      break if target_idx == 0
      parent_idx = parent_idx(target_idx)
    end
  end

  def top_down
    target_idx = 0

    while (has_child?(target_idx))

      a = left_child_idx(target_idx)
      b = right_child_idx(target_idx)

      if @data[b].nil?
        c = a
      else
        c = @data[a] <= @data[b] ? a : b
      end

      if @data[target_idx] > @data[c]
        swap(target_idx, c)
        target_idx = c
      else
        return
      end
    end
  end

  # @param Integer
  # @return Integer
  def left_child_idx(idx)
    (idx * 2) + 1
  end

  # @param Integer
  # @return Integer
  def right_child_idx(idx)
    (idx * 2) + 2
  end

  # @param Integer
  # @return Boolent
  def has_child?(idx)
    ((idx * 2) + 1) < @data.size
  end
end

n, m = gets.split.map(&:to_i)
a = gets.split.map(&:to_i)
e = a.map{|x| -x}
b = PriorityQueue.new(e)
m.times do
  c = b.pop
  b.push(-(-c / 2))
end
ans = 0
while b.size > 0
  ans -= b.pop
end
puts ans

これでもギリギリです。

Python

python.py

import heapq
import math

n, m = map(int, input().split())
a = [-1 * int(i) for i in input().split()]
heapq.heapify(a)
for _ in range(m):
    b = heapq.heappop(a)
    heapq.heappush(a, math.ceil(b / 2))
print(-1 * sum(a))

Pythonのheapqは最小値を取るものですので、マイナスの符号を付けて入れる必要があります。
また、//によるマイナスの割り算は絶対値の大きい方の値を返すので、ここではceilを使用しています。

Java

java.java

import java.util.*;

class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int N = Integer.parseInt(sc.next());
        int M = Integer.parseInt(sc.next());
        PriorityQueue<Long> A = new PriorityQueue<>(Collections.reverseOrder());
        for (int i=0; i<N; i++) {
            A.add(Long.parseLong(sc.next()));
        }
        sc.close();

        for (int i=0; i<M; i++) {
            long new_price = (long)A.poll()/2;
            A.add(new_price);
        }

        long sum = 0;
        for (long a : A) {
            sum += a;
        }
        System.out.println(sum);
    }
}

        PriorityQueue<Long> A = new PriorityQueue<>(Collections.reverseOrder());

Java はreverseOrder()で最大値に対応しています。

||Ruby|Python|Java|
|:--:|:--:|:--:|:--:|:--:|
|コード長 (Byte)|1933|230|673|
|実行時間 (ms)|1981|163|476|
|メモリ (KB)|14004|14536|50024|

まとめ

• ABC 141 D を解いた
• Ruby に詳しくなった
• Python に詳しくなった
• Java に詳しくなった

参照したサイト
Ruby で Priority Queue を実装してみたい
[ruby] Priority Queueの実装