概要

数列 \(A\) の連続する区間 \([l, r]\) において、要素の総和に「区間の長さ \(\times M\)」を加えたコストが \(K\) 以下となる組み合わせの数を求める問題です。

考察

式の変形

まずは、与えられたコストの式を扱いやすい形に変形しましょう。 区間 \([l, r]\) の要素の総和を累積和 \(S_i = A_1 + A_2 + \dots + A_i\) （ただし \(S_0 = 0\)）を用いて表すと、以下のようになります。

\[\text{コスト} = (S_r - S_{l-1}) + (r - l + 1) \times M\]

この式を \(r\) に関する項と \(l-1\) に関する項に整理します。

\[\text{コスト} = (S_r + r \times M) - (S_{l-1} + (l-1) \times M)\]

ここで、\(B_i = S_i + i \times M\) と定義すると、コストは \(B_r - B_{l-1}\) と非常にシンプルな形で表せます。 求める条件は \(\text{コスト} \leq K\) なので、

\[B_r - B_{l-1} \leq K \iff B_{l-1} \geq B_r - K\]

となります。したがって、問題は「\(0 \leq i < j \leq N\) を満たすペア \((i, j)\) のうち、\(B_i \geq B_j - K\) を満たすものの個数を数える」という問題に帰着されます。

高速な数え上げ

\(N\) が最大 \(2 \times 10^5\) であるため、全てのペア \((i, j)\) を調べる \(O(N^2)\) の解法では間に合いません。 \(j\) を \(1\) から \(N\) まで順番に動かしながら、それまでに現れた \(B_i\) （\(0 \leq i < j\)）の中で \(B_i \geq B_j - K\) を満たすものがいくつあるかを高速に数える必要があります。

これは「転倒数」を求める問題と似ており、フェニック木（Binary Indexed Tree, BIT）などのデータ構造を用いることで \(O(\log N)\) で計算可能です。

アルゴリズム

1. 配列 \(B\) の作成: \(B_0 = 0, B_i = B_{i-1} + A_i + M\) として計算します。
2. 座標圧縮: \(B_i\) の値は非常に大きくなる可能性があるため、BIT のインデックスとして直接使えません。出現する \(B_i\) の値をソートし、その「順位（ランク）」に変換する座標圧縮を行います。
3. BIT による数え上げ:
   • \(j\) を \(0\) から \(N\) までループさせます。
   • 各 \(j\) について：
     1. \(B_j - K\) 以上の値が、これまでにいくつ BIT に追加されたかを計算し、答えに加算します。
        • 具体的には、「これまでの総数」から「\(B_j - K\) 未満の個数」を引くことで求められます。
     2. \(B_j\) を BIT に追加します。

計算量

• 時間計算量: \(O(N \log N)\)
  • 累積和の計算に \(O(N)\)、座標圧縮のためのソートに \(O(N \log N)\)、BIT の操作（\(N\) 回）に \(O(N \log N)\) かかります。
• 空間計算量: \(O(N)\)
  • 配列 \(B\) や BIT、座標圧縮用のマップなどを保持するために \(O(N)\) のメモリを使用します。

実装のポイント

• \(B_0\) の扱い: 区間 \([1, r]\) を選ぶ場合は \(l-1 = 0\) となるため、初期値 \(B_0 = 0\) も忘れずに座標圧縮と BIT の対象に含める必要があります。

• 二分探索: \(B_j - K\) が具体的にどのランクに相当するかを判定するために bisect_left を使用します。

• BIT のクエリ: 「\(X\) 以上の個数」を直接求めるのではなく、「全体の個数 - \(X\) 未満の個数」として計算すると実装がスムーズです。

  ソースコード

import sys
import bisect

def solve():
    # Use fast I/O to read all input data at once
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # Parse N, M, K from the first three elements
    N = int(input_data[0])
    M = int(input_data[1])
    K = int(input_data[2])
    
    # Parse the数列 A
    A = list(map(int, input_data[3:3+N]))
    
    # Define B_i = S_i + i * M, where S_i is the prefix sum of A.
    # The condition sum(A_l...A_r) + (r-l+1)*M <= K
    # is equivalent to (S_r + r*M) - (S_{l-1} + (l-1)*M) <= K.
    # Let B_j = S_j + j*M. We need to find the number of pairs (i, j)
    # such that 0 <= i < j <= N and B_j - B_i <= K, which is B_i >= B_j - K.
    B = [0] * (N + 1)
    curr = 0
    for i in range(N):
        curr += A[i] + M
        B[i+1] = curr
        
    # Coordinate compression to use B values as indices in a Binary Indexed Tree (BIT)
    coords = sorted(list(set(B)))
    rank_map = {val: i + 1 for i, val in enumerate(coords)}
    rank_B = [rank_map[val] for val in B]
    size = len(coords)
    
    # Initialize Binary Indexed Tree
    bit = [0] * (size + 1)
    
    ans = 0
    
    # Initial update for B[0] (representing the sum before any elements are chosen)
    # The BIT will store the ranks of B_i values seen so far.
    k = rank_B[0]
    while k <= size:
        bit[k] += 1
        k += k & (-k)
    
    # Main loop to count valid pairs (i, j) where i < j
    for j in range(1, N + 1):
        target = B[j] - K
        # We need to count i < j such that B[i] >= target.
        # bisect_left finds the rank index of the first value in coords >= target.
        idx = bisect.bisect_left(coords, target)
        
        # Count elements in BIT with rank <= idx (i.e., B[i] < target)
        s = 0
        k = idx
        while k > 0:
            s += bit[k]
            k -= k & (-k)
            
        # Total elements currently in BIT is j (B[0], B[1], ..., B[j-1]).
        # Number of B[i] >= target is (total) - (B[i] < target).
        ans += j - s
        
        # Add the rank of current B[j] to the BIT for future iterations
        k = rank_B[j]
        while k <= size:
            bit[k] += 1
            k += k & (-k)
            
    # Output the final count
    sys.stdout.write(str(ans) + '\n')

if __name__ == '__main__':
    solve()

この解説は gemini-3-flash-preview によって生成されました。