C - 区間の合計 / Sum of Intervals Editorial by admin
Qwen3-Coder-480B概要
長さ \(N\) の数列から、連続した部分列(区間)であって、その要素の合計が \(K\) 以下であるものの個数を求めます。
考察
この問題では、すべての連続する区間を全探索して和を計算すると、計算量が \(O(N^2)\) となり、制約 \(N \leq 2 \times 10^5\) では間に合いません。そのため、効率的に条件を満たす区間の数を数える必要があります。
重要な観察として、「区間の和」は累積和を用いることで高速に求められることです。具体的には、数列 \(A\) の累積和 \(prefix\) を定義し、
\[ prefix[i] = A_1 + A_2 + \cdots + A_i \]
とおくと、区間 \([l, r]\) の和は
\[ A_l + A_{l+1} + \cdots + A_r = prefix[r] - prefix[l-1] \]
となります。
今回の問題は「\(prefix[r] - prefix[l-1] \leq K\) となる \((l, r)\) の組の個数」を求めることに他なりません。これを変形すると
\[ prefix[l-1] \geq prefix[r] - K \]
となります。つまり、各右端 \(r\) について、\(prefix[0], prefix[1], \ldots, prefix[r-1]\) の中で、\(prefix[r] - K\) 以上となるものの個数を数えれば良いです。
このとき、累積和は非減少列ではないですが、探索範囲が左側に限定されているため、毎回の探索範囲内でソートされた状態を保つことはできません。しかし、累積和の添字順に沿って走査しているので、二分探索が有効になります。bisect_left を使うことで、条件を満たす最小のインデックスを高速に求められます。
アルゴリズム
- 入力から数列 \(A\) を受け取り、累積和配列
prefixを構築します。 - 各右端 \(r = 1, 2, \ldots, N\) に対して以下を行います:
- \(target = prefix[r] - K\) を計算します。
- 配列
prefix[0..r-1]の中で、値が \(target\) 以上となる最小のインデックス \(pos\) を二分探索で求めます。 - このとき、\(l-1\) が \(pos, pos+1, \ldots, r-1\) となるような \(l\) は全部で \(r - pos\) 個あります。
- この数を答えに加算します。
- 最終的な答えを出力します。
計算量
- 時間計算量: \(O(N \log N)\)
(各 \(r\) に対して二分探索を行うため) - 空間計算量: \(O(N)\)
(累積和配列を保持するため)
実装のポイント
- 累積和のサイズは \(N+1\) にして、
prefix[0] = 0としておくと、区間の左端が 1 から始まる場合でも自然に扱えます。 - Python の
bisect_leftを使用する際、探索範囲を0からr(つまりprefix[:r])に限定する必要があります。これにより、正しい範囲での二分探索が可能です。
## ソースコード
```python
import sys
from bisect import bisect_left
def main():
input = sys.stdin.read
data = input().split()
N = int(data[0])
K = int(data[1])
A = list(map(int, data[2:]))
# 累積和を計算
prefix = [0] * (N + 1)
for i in range(N):
prefix[i + 1] = prefix[i] + A[i]
count = 0
# 各右端 r に対して、条件を満たす左端 l の個数を数える
for r in range(1, N + 1):
# 区間 [l, r] の和 = prefix[r] - prefix[l-1] <= K
# つまり prefix[l-1] >= prefix[r] - K
target = prefix[r] - K
# prefix[0..r-1] の中で、prefix[l-1] >= target となる最小の l-1 を見つける
# これは prefix[0..r-1] が非減少なので二分探索可能
pos = bisect_left(prefix, target, 0, r) # prefix[0..r-1] の範囲で探索
count += r - pos # l-1 が pos, pos+1, ..., r-1 に対応する l が存在
print(count)
if __name__ == "__main__":
main()
この解説は qwen3-coder-480b によって生成されました。