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

はじめに

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 --- 配列二分法アルゴリズム