Official

C - 区間の合計 / Sum of Intervals Editorial by admin Editorial for printing

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 を使うことで、条件を満たす最小のインデックスを高速に求められます。

アルゴリズム

  1. 入力から数列 \(A\) を受け取り、累積和配列 prefix を構築します。
  2. 各右端 \(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\) 個あります。
    • この数を答えに加算します。
  3. 最終的な答えを出力します。

計算量

  • 時間計算量: \(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 によって生成されました。

posted:
last update: