Ruby と Python で解く AtCoder ABC172 C 累積和 二分探索

17122 ワード
Ruby AtCoder Python 新人プログラマ応援 Ruby テキストリンク

はじめに

AtCoder Beginner Contest 172 に参加しました。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

AtCoder Beginner Contest C - Tsundoku
Difficulty: 878

今回のテーマ、累積和 + 二分探索

最初は、動的計画法と考え解いていましたがTLEを貰い大幅な時間のロス、二分探索に切り替えました。
入力例1ですと、次の表の各行で、机Aと机Bの累積和がK分を超えない本数の最大値を求めます。

1 2 3 4 5 6 7
60 90 120 80 150 80 150
60 90 80 150 80 150
60 80 150 80 150
80 150 80 150

Ruby

ruby.rb
n, m, k = gets.split.map(&:to_i)
a = gets.split.map(&:to_i)
b = gets.split.map(&:to_i)
c = Array.new(n + 1, 0)
n.times do |i|
  c[i + 1] = c[i] + a[i]
end
d = Array.new(m + 1, 0)
m.times do |i|
  d[i + 1] = d[i] + b[i]
end
cnt = 0
0.upto(n) do |i|
  break if c[i] > k
  e = d.bsearch_index{_1 > k - c[i]}
  if e.nil?
    cnt = i + m if cnt < i + m
  else
    cnt = i + e - 1 if cnt < i + e - 1
  end
end
puts cnt
ruby.rb
c = Array.new(n + 1, 0)
n.times do |i|
  c[i + 1] = c[i] + a[i]
end
d = Array.new(m + 1, 0)
m.times do |i|
  d[i + 1] = d[i] + b[i]
end

累積和です。

ruby.rb
  e = d.bsearch_index{_1 > k - c[i]}
  if e.nil?
    cnt = i + m if cnt < i + m
  else
    cnt = i + e - 1 if cnt < i + e - 1
  end

二分探索です。配列の最大値を超えた場合、nilを返しますので、その処理が入っています。

Python

python.py
from sys import stdin
import bisect

def main():
    input = stdin.readline
    n, m, k = map(int, input().split())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))
    c = [0] * (n + 1)
    for i in range(n):
        c[i + 1] = c[i] + a[i]
    d = [0] * (m + 1)
    for i in range(m):
        d[i + 1] = d[i] + b[i]
    cnt = 0
    for i in range(n + 1):
        if c[i] > k:
            break
        e = bisect.bisect_right(d, k - c[i])
        if cnt < i + e - 1:
            cnt = i + e - 1
    print(cnt)
main()
python.py
        e = bisect.bisect_right(d, k - c[i])
        if cnt < i + e - 1:
            cnt = i + e - 1

二分探索です。配列の最大値を超えた場合、配列の要素数+1を返します。

Ruby Python
コード長 (Byte) 433 570
実行時間 (ms) 349 246
メモリ (KB) 44860 47636

まとめ

  • ABC 172 C を解いた
  • Ruby に詳しくなった
  • Python に詳しくなった

参照したサイト
bisect --- 配列二分法アルゴリズム

pythonのCounter.most_common()
『C++言語基礎』プログラム読書——ポインタ、const、static(二)