AtCoder ABC 130 D - Enough Array (Python) - あるエンジニアのAtCoder奮闘記

irisruneです。AtCoderで言語アップデートがあるそうで、色々な言語が使いやすくなるといいですね。

問題

連続部分列とみると累積和を取りたくなりませんか？累積和を取った場合は二分探索が手っ取り早いと思います。

import sys
import bisect
import numpy as np


def main():
    n, k = [int(e) for e in input().split()]
    a = [int(e) for e in input().split()]
    cumsum = [0]
    for i in range(n):
        cumsum.append(cumsum[i] + a[i])
    cumsum.sort()
    ans = 0
    for cs in cumsum:
        if cs < k:
            continue
        ans += bisect.bisect(cumsum, cs - k)
    print(ans)


main()

累積和を用いても連続部分列の和を全通り求めるのに計算量が $O(N^2)$ かかってしまいます。しかし累積和の性質を使うことで、二分探索を用いて解くことができます。

詳しく説明すると、累積和 $C_i$ について $C_i-K\geq C_j$ となる場合、 $C_i-C_j\geq K$ となります。これは累積和の性質よりある1つの連続部分列の和が $K$ 以上であることを示します（元の数列の各要素が負でないことが前提です）。そのため、累積和 $C_i$ をソートし、 $C_i-K\geq C_j$ となる組の数を二分探索を用いて求めることで、全体の計算量を $O(N\log N)$ で抑えることができます。

公式解説では尺取り法を勧められていますが、累積和が思いついて、かつライブラリが使える前提であれば二分探索でも楽に実装できると思います。

雑記

C問題は解く前に本質的な画像が流れてきてしまって全く自力でないACをしました。
PythonでTLEした問題をPyPyなら解けたという事例が今のところないので、PyPyの使いどころがわからないのが悩みです。
- この記事の問題など、PyPy推奨らしいのですが自分のコードでは通りませんでした。
- 記事にはしていないですが、逆にPythonで通るのにPyPyで通らないケースがあったのでPythonの方が信頼できてしまうんですよね…