AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC130 D - Enough Array

問題原文

atcoder.jp

問題要旨

長さ  N の正整数列  A の連続する部分列であって、その総和が  K 以上となるものの個数を求めよ。
ただし、部分列が同じであっても取り出した位置が異なるなら別のものとして数えることとする。

  •  1 \leq N \leq 10^{5}
  •  1 \leq K \leq 10^{10}
  •  1 \leq a_i \leq 10^{5}

解法

 1 \leq a_i より、部分列の総和には単調性があることがわかる。
つまり、 \sum_{i=s}^{t} a_i \geq K のとき、  u > t となる全ての  u において、 \sum_{i=s}^{u}a_i > Kが必ず成り立つ。 よって、そのような  u に関しては、部分列の総和を計算する必要はない。
逆に言えば、ある  s について、初めて  \sum_{i=s}^{t} a_i \geq K となるような  t がわかれば、その  s について、「いくつ条件を満たすような部分列の終点が存在するか」を計算できる。

f:id:AT274:20200102140015p:plain

で、そういう「 K 以上のになるかどうかの境目」を見つけるにはいくつか方法がある。

累積和 + 二分探索

累積和 + 二分探索で、各  s について境目を高速に求めることができる。
K以上が条件なので、 bisect_left を使う点に少し注意。

f:id:AT274:20200102143733p:plain
累積和 + 二分探索のイメージ

ある  s を始点とする部分列を考える時、 s - 1 までの和は使ってはいけないので引く必要がある。
→引くのではなく、 K に加算することで条件達成を厳しくしている、ようなイメージ。

計算量は  O(NlogN)

実装
from itertools import accumulate
from bisect import bisect_left

N, K = map(int, input().split())
A = list(map(int, input().split()))

A = [0] + list(accumulate(A))
ans = 0
for i in range(1, N + 1):
    ans += N - bisect_left(A, A[i - 1] + K) + 1

print(ans)

しゃくとり法

しゃくとり法も使える。どちらかというとこちらのがイメージするのは簡単かも。 ある s を始点とする部分列を考えていって、  K 以上となった地点でストップするイメージ。 最後まで見ても和が  K 以上にならないパターンがあることに少し注意する。

計算量は  O(N)

実装
N, K = map(int, input().split())
A = list(map(int, input().split()))

ans, r, s = 0, 0, 0
for l in range(N):
    while (r < N) and (s < K):
        s += A[r]
        r += 1

    if s >= K:
        ans += N - r + 1
    s -= A[l]

print(ans)

感想

前はにぶたん解のが得意だったけど今見たらしゃくとり解のがイメージしやすかった。 どちらでも瞬殺できるようになっておきたい。
単調性がないか? 本当に単調性があるのか? は常に意識したい。(今回だと  a_i が負になりうるなら、単調性は成り立たない)