ABC130 D - Enough Array
問題原文
問題要旨
長さ の正整数列 の連続する部分列であって、その総和が 以上となるものの個数を求めよ。
ただし、部分列が同じであっても取り出した位置が異なるなら別のものとして数えることとする。
解法
より、部分列の総和には単調性があることがわかる。
つまり、 のとき、 となる全ての において、が必ず成り立つ。
よって、そのような に関しては、部分列の総和を計算する必要はない。
逆に言えば、ある について、初めて となるような がわかれば、その について、「いくつ条件を満たすような部分列の終点が存在するか」を計算できる。
で、そういう「 以上のになるかどうかの境目」を見つけるにはいくつか方法がある。
累積和 + 二分探索
累積和 + 二分探索で、各 について境目を高速に求めることができる。
K以上
が条件なので、 bisect_left
を使う点に少し注意。
ある を始点とする部分列を考える時、 までの和は使ってはいけないので引く必要がある。
→引くのではなく、 に加算することで条件達成を厳しくしている、ようなイメージ。
計算量は 。
実装
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)
しゃくとり法
しゃくとり法も使える。どちらかというとこちらのがイメージするのは簡単かも。 ある を始点とする部分列を考えていって、 以上となった地点でストップするイメージ。 最後まで見ても和が 以上にならないパターンがあることに少し注意する。
計算量は 。
実装
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)
感想
前はにぶたん解のが得意だったけど今見たらしゃくとり解のがイメージしやすかった。 どちらでも瞬殺できるようになっておきたい。
単調性がないか?
本当に単調性があるのか?
は常に意識したい。(今回だと が負になりうるなら、単調性は成り立たない)