Ruby と crystal で解く AtCoder ABC253 C

はじめに

AtCoder さん、ありがとうございます。

C - Max - Min Query

出題の趣旨として、配列のMaxやMin(やSort等)は毎回計算するとコストがかかるので、TLEを避ける工夫が必要です。

Hash + メモ

Ruby

q = gets.to_i
h = Hash.new(0)
h_max = -1
h_min = 1000000001
q.times do
  s = gets.chomp
  case s[0]
  when '1'
    x = s.split[1].to_i
    h[x] += 1
    h_max = x if h_max < x
    h_min = x if h_min > x
  when '2'
    _, x, c = s.split.map(&:to_i)
    if h[x] > c
      h[x] -= c
    else
      h.delete(x)
      h_max = h.empty? ? -1 : h.keys.max if h_max == x
      h_min = h.empty? ? 1000000001 : h.keys.min if h_min == x
    end
  when '3'
    puts h_max - h_min
  end
end

h_max/h_minで最大値/最小値を記録しておき、その値が削除される場合のみ更新することにより、計算コストを削減しています。
ハッシュは個数を管理しています。

crystal

q = read_line.to_i
h = Hash(Int32, Int32).new(0)    
h_max = -1
h_min = 1000000001
q.times do
  s = read_line
  case s[0]
  when '1'
    x = s.split[1].to_i
    h[x] += 1
    h_max = x if h_max < x       
    h_min = x if h_min > x       
  when '2'
    _, x, c = s.split.map(&.to_i)
    if h[x] > c
      h[x] -= c
    else
      h.delete(x)
      h_max = h.empty? ? -1 : h.keys.max if h_max == x
      h_min = h.empty? ? 1000000001 : h.keys.min if h_min == x
    end
  when '3'
    puts h_max - h_min
  end
end

この問題では、最大値と最小値以外の中間の要素について順番を気にする必要はないので、Hashの様なシンプルなデータ構造で対応可能でした。

	Ruby	crystal
実行時間(ms)	309	79

PriorityQueue

ステップアップのため、有名なデータ構造も押さえておきます。

ユーザ解説では、平方分割・heapq・Binary Indexed Tree が紹介されています。

最大値（最小値）で計算コストを抑えるとなると優先度付きキュー。
Pythonならばheapqがありますが、Rubyやcrystalでは自作で対応する必要があります。

Ruby

# Priority Queue
# Reference: https://github.com/python/cpython/blob/master/Lib/heapq.py
class PriorityQueue
  # By default, the priority queue returns the maximum element first.
  # If a block is given, the priority between the elements is determined with it.
  # For example, the following block is given, the priority queue returns the minimum element first.
  # `PriorityQueue.new { |x, y| x < y }`
  #
  # A heap is an array for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for all k, counting elements from 0.
  def initialize(array = [], &comp)
    @heap = array
    @comp = comp || proc { |x, y| x > y }
    heapify
  end

  attr_reader :heap

  # Push new element to the heap.
  def push(item)
    shift_down(0, @heap.push(item).size - 1)
  end

  alias << push
  alias append push

  # Pop the element with the highest priority.
  def pop
    latest = @heap.pop
    return latest if empty?

    ret_item = heap[0]
    heap[0] = latest
    shift_up(0)
    ret_item
  end

  # Get the element with the highest priority.
  def get
    @heap[0]
  end

  alias top get

  # Returns true if the heap is empty.
  def empty?
    @heap.empty?
  end

  private

  def heapify
    (@heap.size / 2 - 1).downto(0) { |i| shift_up(i) }
  end

  def shift_up(pos)
    end_pos = @heap.size
    start_pos = pos
    new_item = @heap[pos]
    left_child_pos = 2 * pos + 1

    while left_child_pos < end_pos
      right_child_pos = left_child_pos + 1
      if right_child_pos < end_pos && @comp.call(@heap[right_child_pos], @heap[left_child_pos])
        left_child_pos = right_child_pos
      end
      # Move the higher priority child up.
      @heap[pos] = @heap[left_child_pos]
      pos = left_child_pos
      left_child_pos = 2 * pos + 1
    end
    @heap[pos] = new_item
    shift_down(start_pos, pos)
  end

  def shift_down(star_pos, pos)
    new_item = @heap[pos]
    while pos > star_pos
      parent_pos = (pos - 1) >> 1
      parent = @heap[parent_pos]
      break if @comp.call(parent, new_item)

      @heap[pos] = parent
      pos = parent_pos
    end
    @heap[pos] = new_item
  end
end

q = gets.to_i
h_max = PriorityQueue.new
h_min = PriorityQueue.new { |x, y| x < y }
h = Hash.new(0)
q.times do
  s = gets.chomp
  case s[0]
  when '1'
    x = s.split[1].to_i
    if h[x] == 0
      h_max << x
      h_min << x
    end
    h[x] += 1
  when '2'
    _, x, c = s.split.map(&:to_i)
    if h[x] >= c
      h[x] -= c
    else
      h[x] = 0
    end
  when '3'
    while h[h_max.get] == 0
      h_max.pop
    end
    while h[h_min.get] == 0
      h_min.pop
    end
    puts h_max.get - h_min.get
  end
end

ご存知、Rubyのatcoder用libraryからPriorityQueueをお借りしています。
ここでは、最大値用のキューと最小値用のキューを用意しています。

crystal

# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

module AtCoder
  # Implements standard priority queue like [std::priority_queue](https://en.cppreference.com/w/cpp/container/priority_queue).
  #
  # ```
  # q = AtCoder::PriorityQueue(Int64).new
  # q << 1_i64
  # q << 3_i64
  # q << 2_i64
  # q.pop # => 3
  # q.pop # => 2
  # q.pop # => 1
  # ```
  class PriorityQueue(T)
    getter heap : Array(T)

    def initialize
      initialize {|a, b| a <= b}
    end

    # Initializes queue with the custom comperator.
    #
    # If the second argument `b` should be popped earlier than
    # the first argument `a`, return `true`. Else, return `false`.
    #
    # ```
    # q = AtCoder::PriorityQueue(Int64).new {|a, b| a >= b}
    # q << 1_i64
    # q << 3_i64
    # q << 2_i64
    # q.pop # => 1
    # q.pop # => 2
    # q.pop # => 3
    # ```
    def initialize(&block : T, T -> Bool)
      @heap = Array(T).new
      @compare_proc = block
    end

    # Pushes value into the queue.
    def push(v : T)
      @heap << v
      index = @heap.size - 1
      while index != 0
        parent = (index - 1) // 2
        if @compare_proc.call(@heap[index], @heap[parent])
          break
        end
        @heap[parent], @heap[index] = @heap[index], @heap[parent]
        index = parent
      end
    end

    # Alias of `push`
    def <<(v : T)
      push(v)
    end

    # Pops value from the queue.
    def pop
      if @heap.size == 0
        return nil
      end
      if @heap.size == 1
        return @heap.pop
      end
      ret = @heap.first
      @heap[0] = @heap.pop
      index = 0
      while index * 2 + 1 < @heap.size
        child = if index * 2 + 2 < @heap.size && !@compare_proc.call(@heap[index * 2 + 2], @heap[index * 2 + 1])
          index * 2 + 2
        else
          index * 2 + 1
        end
        if @compare_proc.call(@heap[child], @heap[index])
          break
        end
        @heap[child], @heap[index] = @heap[index], @heap[child]
        index = child
      end
      ret
    end

    # Returns `true` if the queue is empty.
    delegate :empty?, to: @heap

    # Returns size of the queue.
    delegate :size, to: @heap
  end
end

q = read_line.to_i
h_max = AtCoder::PriorityQueue(Int32).new
h_min = AtCoder::PriorityQueue(Int32).new {|a, b| a >= b}
h = Hash(Int32, Int32).new(0)
q.times do
  s = read_line
  case s[0]
  when '1'
    x = s.split[1].to_i
    if h[x] == 0
      h_max << x
      h_min << x
    end
    h[x] += 1
  when '2'
    _, x, c = s.split.map(&.to_i)
    if h[x] >= c
      h[x] -= c
    else
      h[x] = 0
    end
  when '3'
    t_max = h_max.pop
    while h[t_max] == 0
      t_max = h_max.pop
    end
    t_min = h_min.pop
    while h[t_min] == 0
      t_min = h_min.pop
    end
    puts t_max.not_nil! - t_min.not_nil!
    h_max << t_max.not_nil!
    h_min << t_min.not_nil!
  end
end

こちらからPriorityQueueをお借りしました。
getsメソッドがないので、popしてpushしています。

	Ruby	crystal
実行時間(ms)	360	109

Fenwick Tree

学習中

まとめ

require "PriorityQueue"と記載しておくと提出時に、PriorityQueueのコードをコピペするような仕組みがあると便利でしょうか。