概要

連続する店舗の売上合計が \(K\) 円以上となる区間 \((l, r)\) の個数を、累積和と二分探索を用いて効率的に数え上げる問題です。

考察

素朴なアプローチとその問題点

まず思いつくのは、すべての \((l, r)\) の組み合わせを試して、各区間の合計を計算する方法です。

for l in 1 to N:
    for r in l to N:
        区間[l, r]の合計を計算してKと比較

しかし、この方法では \((l, r)\) の組み合わせが約 \(\frac{N^2}{2}\) 通りあり、\(N = 2 \times 10^5\) の場合は約 \(2 \times 10^{10}\) 回の計算が必要となり、TLE（時間制限超過）になります。

重要な気づき

1. 累積和を使えば区間の合計を \(O(1)\) で求められる

   • \(\text{prefix}[i] = V_1 + V_2 + \cdots + V_{i}\) と定義すると
   • \(V_l + V_{l+1} + \cdots + V_r = \text{prefix}[r] - \text{prefix}[l-1]\)

2. すべての \(V_i \geq 1\) なので、累積和は単調増加

   • \(\text{prefix}[0] < \text{prefix}[1] < \text{prefix}[2] < \cdots < \text{prefix}[N]\)
   • 単調増加な配列では二分探索が使えます！

3. 条件の言い換え

   • \(V_l + \cdots + V_r \geq K\)
   • \(\Leftrightarrow \text{prefix}[r] - \text{prefix}[l-1] \geq K\)
   • \(\Leftrightarrow \text{prefix}[r] \geq \text{prefix}[l-1] + K\)

つまり、各 \(l\) に対して「\(\text{prefix}[r] \geq \text{prefix}[l-1] + K\) を満たす最小の \(r\)」を二分探索で見つければ、そこから \(N\) までのすべての \(r\) が条件を満たします。

アルゴリズム

1. 累積和の計算

   • \(\text{prefix}[0] = 0\)
   • \(\text{prefix}[i] = \text{prefix}[i-1] + V_i\) （\(i = 1, 2, \ldots, N\)）

2. 各 \(l\) について二分探索

   • \(\text{target} = \text{prefix}[l-1] + K\) を計算
   • 累積和配列から \(\text{target}\) 以上となる最小のインデックス \(m\) を二分探索で求める
   • \(m \leq N\) かつ \(m \geq l\) ならば、\(r = m, m+1, \ldots, N\) の \((N - m + 1)\) 個が条件を満たす

3. すべての \(l\) について合計を求める

具体例

\(N = 4\), \(K = 6\), \(V = [2, 3, 4, 1]\) の場合： - \(\text{prefix} = [0, 2, 5, 9, 10]\)

\(l = 1\) のとき： - \(\text{target} = 0 + 6 = 6\) - \(\text{prefix}[r] \geq 6\) となる最小の \(r\) は \(3\)（\(\text{prefix}[3] = 9\)） - \(r = 3, 4\) の2通りが条件を満たす

計算量

• 時間計算量: \(O(N \log N)\)
  • 累積和の計算に \(O(N)\)
  • 各 \(l\) について二分探索を行うので \(O(N \log N)\)
• 空間計算量: \(O(N)\)
  • 累積和配列の保存に \(O(N)\)

実装のポイント

• 累積和のインデックス管理: コード中では0-indexedで実装しているため、prefix[i+1] = prefix[i] + V[i] として prefix[0] = 0 としています。区間 \([l, r]\)（0-indexed）の合計は prefix[r+1] - prefix[l] です。

• 二分探索の境界条件: bisect_left で見つけたインデックスが \(l+1\) 以上であることを確認する必要があります（\(r \geq l\) の制約があるため）。

• オーバーフロー対策: \(K\) が最大 \(10^{14}\)、\(V_i\) が最大 \(10^9\)、\(N\) が最大 \(2 \times 10^5\) なので、累積和は最大で約 \(2 \times 10^{14}\) になります。Pythonでは整数のオーバーフローを気にする必要はありません。

  ソースコード

import bisect

def solve():
    N, K = map(int, input().split())
    V = list(map(int, input().split()))
    
    # 累積和を計算
    # prefix[i] = V[0] + V[1] + ... + V[i-1]
    # prefix[0] = 0
    prefix = [0] * (N + 1)
    for i in range(N):
        prefix[i + 1] = prefix[i] + V[i]
    
    # V[l] + V[l+1] + ... + V[r] = prefix[r+1] - prefix[l]
    # これが >= K となる条件は prefix[r+1] - prefix[l] >= K
    # つまり prefix[r+1] >= prefix[l] + K
    
    # すべてのV_iが正なので、prefixは単調増加
    # 各lに対して、prefix[r+1] >= prefix[l] + K となる最小のr+1を二分探索で見つける
    # そのようなr+1がm以上N+1以下であれば、(N+1) - m 個の組が条件を満たす
    
    count = 0
    for l in range(N):
        # prefix[l] + K 以上となる最小のインデックスを探す
        target = prefix[l] + K
        # bisect_leftで target 以上となる最初の位置を探す
        idx = bisect.bisect_left(prefix, target)
        
        # idx は l+1 以上 N+1 以下でなければならない（r >= l なので r+1 >= l+1）
        if idx <= N:
            # r+1 が idx から N までの範囲で条件を満たす
            # r は idx-1 から N-1 まで
            # ただし r >= l なので idx-1 >= l, つまり idx >= l+1
            start = max(idx, l + 1)
            if start <= N:
                count += N - start + 1
    
    print(count)

solve()

この解説は claude4.5opus によって生成されました。