Ruby と Python で解く AtCoder ABC130 D 累積和 二分探索

はじめに

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

今回のお題

AtCoder Beginner Contest D - Enough Array
Difficulty: 859

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

問題文(条件)連続部分列に含まれる全ての要素の値の和は、K 以上である。より累積和を使用することが分かりますが、1≦N≦100_000ですので2重にループを回すとTLEになります。
そこで、累積和は単調増加ですので、計算量を抑えるために二分探索も併用します。

Ruby

ruby.rb

n, k = gets.split.map(&:to_i)
a = gets.split.map(&:to_i)
acc = 0
x = [0] + a.map{|v| acc += v}
cnt = 0
x.each do |y|
  j = x.bsearch_index{|z| z - y >= k}
  if j == nil
    break
  else
    cnt += n - j + 1
  end
end
puts cnt

accumulate.rb

acc = 0
x = [0] + a.map{|v| acc += v}

累積和はここで計算しています。
追記
コメントでいただいたコードに修正しました。

bsearch.rb

  j = x.bsearch_index{|z| z - y >= k}

二分探索ですが、C++ でいうところのlower_boundが、Ruby ではbsearch若しくはbsearch_indexになります。

Python

python.py

from sys import stdin

def main():
    import bisect
    import itertools

    input = stdin.readline
    n, k = map(int, input().split())
    a = list(map(int, input().split()))
    x = [0] + list(itertools.accumulate(a))
    cnt = 0
    for z in x:
        j = bisect.bisect_left(x, k + z)
        if j <= n:
            cnt += n - j + 1
    print(cnt)
main()

accumulate.py

    x = [0] + list(itertools.accumulate(a))

Python では、累積和を求めるaccumulateという関数があります。

bisect.py

        j = bisect.bisect_left(x, k + z)

二分探索ですが、bisect若しくはbisect_left bisect_rightを使用します。
探索する数値が配列に含まれる場合、bisectはその数値の右側、bisect_leftは左側の位置を返します。

for.py

    for z in x:
    for i in range(n):

range()でループを回すとTLEしました。

	Ruby	Python
コード長 (Byte)	238	377
実行時間 (ms)	165	101
メモリ (KB)	11780	14196

まとめ

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

参照したサイト

instance method Array#bsearch
すごいぞitertoolsくん
bisect --- 配列二分法アルゴリズム