Overview

This problem asks us to find the number of contiguous subarray intervals \([l, r]\) of sequence \(A\) where the cost — defined as the sum of elements plus “the length of the interval \(\times M\)” — is at most \(K\).

Analysis

Transforming the Expression

First, let’s transform the given cost expression into a more manageable form. Expressing the sum of elements in interval \([l, r]\) using prefix sums \(S_i = A_1 + A_2 + \dots + A_i\) (where \(S_0 = 0\)), we get:

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

Rearranging this into terms involving \(r\) and terms involving \(l-1\):

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

Defining \(B_i = S_i + i \times M\), the cost can be expressed very simply as \(B_r - B_{l-1}\). Since the condition we want is \(\text{Cost} \leq K\):

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

Therefore, the problem reduces to “counting the number of pairs \((i, j)\) satisfying \(0 \leq i < j \leq N\) and \(B_i \geq B_j - K\).”

Efficient Counting

Since \(N\) can be up to \(2 \times 10^5\), an \(O(N^2)\) approach that checks all pairs \((i, j)\) is too slow. As we iterate \(j\) from \(1\) to \(N\), we need to efficiently count how many of the previously seen \(B_i\) values (\(0 \leq i < j\)) satisfy \(B_i \geq B_j - K\).

This is similar to counting inversions, and can be computed in \(O(\log N)\) per query using data structures such as a Fenwick Tree (Binary Indexed Tree, BIT).

Algorithm

1. Construct array \(B\): Compute as \(B_0 = 0, B_i = B_{i-1} + A_i + M\).
2. Coordinate compression: Since \(B_i\) values can be very large, they cannot be used directly as BIT indices. Sort the occurring \(B_i\) values and convert them to their “ranks” via coordinate compression.
3. Counting with BIT:
   • Loop \(j\) from \(0\) to \(N\).
   • For each \(j\):
     1. Compute how many values greater than or equal to \(B_j - K\) have been added to the BIT so far, and add this count to the answer.
        • Specifically, this is obtained by subtracting “the count of values less than \(B_j - K\)” from “the total count so far.”
     2. Add \(B_j\) to the BIT.

Complexity

• Time complexity: \(O(N \log N)\)
  • Computing prefix sums takes \(O(N)\), sorting for coordinate compression takes \(O(N \log N)\), and BIT operations (\(N\) times) take \(O(N \log N)\).
• Space complexity: \(O(N)\)
  • \(O(N)\) memory is used to store array \(B\), the BIT, the coordinate compression map, etc.

Implementation Notes

• Handling \(B_0\): When selecting interval \([1, r]\), we have \(l-1 = 0\), so the initial value \(B_0 = 0\) must be included in both the coordinate compression and BIT operations.

• Binary search: Use bisect_left to determine which rank \(B_j - K\) corresponds to.

• BIT queries: Rather than directly computing “the count of values \(\geq X\)”, it is smoother to implement it as “total count - count of values \(< X\)”.

  Source Code

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()

This editorial was generated by gemini-3-flash-preview.