Overview

This problem asks us to count the number of intervals \((l, r)\) where the “sum of elements + interval length × M” cost is at most \(K\) when selecting a contiguous interval from a sequence.

Analysis

Rephrasing the Problem

The cost of interval \([l, r]\) is: $\(\text{cost} = \sum_{i=l}^{r} A_i + (r - l + 1) \times M\)$

If we define a new sequence \(B_i = A_i + M\): $\(\text{cost} = \sum_{i=l}^{r} (A_i + M) = \sum_{i=l}^{r} B_i\)$

Thus, we can reduce this to “counting intervals where the contiguous subarray sum of \(B\) is at most \(K\)”.

Transformation Using Prefix Sums

Let \(S\) be the prefix sum of \(B\) (\(S_0 = 0\), \(S_i = B_0 + B_1 + \cdots + B_{i-1}\)).

The cost of interval \([l, r]\) can be expressed as \(S_{r+1} - S_l\), so the condition we seek is: $\(S_{r+1} - S_l \leq K\)$

Rearranging this: $\(S_l \geq S_{r+1} - K\)$

Problem with the Naive Approach

For each \(r\), exhaustively searching for all \(l\) satisfying the condition takes \(O(N^2)\) time, which causes TLE when \(N \leq 2 \times 10^5\).

Solution

Process each \(r\) in order, and efficiently find “the count of \(S_l\) values (\(l = 0, 1, \ldots, r\)) seen so far that satisfy \(S_l \geq S_{r+1} - K\)”.

This can be done in \(O(\log N)\) per query using coordinate compression + Binary Indexed Tree (BIT).

Algorithm

1. Preprocessing: Compute \(B_i = A_i + M\) and calculate the prefix sum \(S\)
2. Coordinate Compression: Sort \(S_0, S_1, \ldots, S_N\) and each threshold \(S_{r+1} - K\) together, converting values to small integers
3. Counting with BIT:
   • Initially add \(S_0\) to the BIT
   • For \(r = 0, 1, \ldots, N-1\):
     • Calculate threshold \(\text{threshold} = S_{r+1} - K\)
     • Query the BIT for “count of elements with value ≥ threshold” and add to the answer
     • Add \(S_{r+1}\) to the BIT (can be used as \(l = r+1\) for the next \(r\))

Concrete Example

For \(N=3\), \(M=2\), \(K=10\), \(A = [1, 2, 3]\): - \(B = [3, 4, 5]\) - \(S = [0, 3, 7, 12]\)

When \(r=0\): threshold is \(S_1 - K = 3 - 10 = -7\). Since \(S_0 = 0 \geq -7\), there is 1 valid interval. When \(r=1\): threshold is \(S_2 - K = 7 - 10 = -3\). Both \(S_0 = 0\) and \(S_1 = 3\) are \(\geq -3\), so 2 valid intervals. When \(r=2\): threshold is \(S_3 - K = 12 - 10 = 2\). \(S_0 = 0 < 2\), \(S_1 = 3 \geq 2\), \(S_2 = 7 \geq 2\), so 2 valid intervals.

Total: \(1 + 2 + 2 = 5\) intervals.

Complexity

• Time complexity: \(O(N \log N)\) (sorting for coordinate compression and BIT operations for each \(r\))
• Space complexity: \(O(N)\) (prefix sum, coordinate compression array, BIT)

Implementation Notes

• Coordinate Compression: Both the values of \(S\) and thresholds \(S_{r+1} - K\) must be included in the compression
• 1-indexed BIT: BIT is typically implemented as 1-indexed, so be careful with index conversion
• Overflow Prevention: Since \(K\) can be up to \(10^{18}\), Python handles this fine, but other languages need 64-bit integers
• Range Queries: “Count of values ≥ threshold” is calculated as “total count - count of values < threshold”

Source Code

import sys
from bisect import bisect_right

def solve():
    input_data = sys.stdin.read().split()
    N = int(input_data[0])
    M = int(input_data[1])
    K = int(input_data[2])
    A = [int(input_data[3 + i]) for i in range(N)]
    
    # Define B[i] = A[i] + M, so the cost of interval [l,r] becomes sum(B[l..r])
    # Cost = sum(A[l..r]) + (r-l+1)*M = sum(B[l..r])
    B = [A[i] + M for i in range(N)]
    
    # Prefix sum S[0] = 0, S[i] = B[0] + B[1] + ... + B[i-1]
    # Cost of interval [l,r] = S[r+1] - S[l]
    # Count (l, r) pairs satisfying S[r+1] - S[l] <= K (0 <= l <= r < N)
    
    S = [0] * (N + 1)
    for i in range(N):
        S[i + 1] = S[i] + B[i]
    
    # For each r, count l values where S[r+1] - S[l] <= K
    # Count l (0 <= l <= r) where S[l] >= S[r+1] - K
    
    # Use segment tree or BIT to index by S[l] values
    # and count elements where S[l] >= threshold
    
    # Coordinate compression
    # Required values: S[0], S[1], ..., S[N] and each S[r+1] - K
    all_values = []
    for i in range(N + 1):
        all_values.append(S[i])
    thresholds = []
    for r in range(N):
        thresholds.append(S[r + 1] - K)
        all_values.append(S[r + 1] - K)
    
    # Coordinate compression
    sorted_values = sorted(set(all_values))
    value_to_idx = {v: i for i, v in enumerate(sorted_values)}
    num_values = len(sorted_values)
    
    # BIT (Binary Indexed Tree) for counting
    # BIT uses 1-indexed
    bit = [0] * (num_values + 2)
    
    def bit_add(i, delta=1):
        i += 1  # 1-indexed
        while i <= num_values + 1:
            bit[i] += delta
            i += i & (-i)
    
    def bit_sum(i):
        # Sum of [0, i] (0-indexed up to i)
        i += 1  # 1-indexed
        s = 0
        while i > 0:
            s += bit[i]
            i -= i & (-i)
        return s
    
    def bit_range_sum(l, r):
        # Sum of [l, r] (0-indexed)
        if l > r:
            return 0
        if l == 0:
            return bit_sum(r)
        return bit_sum(r) - bit_sum(l - 1)
    
    # Count l values where S[l] >= threshold
    # This is the count of elements with coordinate-compressed index >= threshold's index
    
    ans = 0
    
    # Add l = 0 first
    bit_add(value_to_idx[S[0]])
    
    for r in range(N):
        # With r fixed, l ranges from 0 to r
        # Count l values where S[l] >= S[r+1] - K
        threshold = thresholds[r]
        threshold_idx = value_to_idx[threshold]
        
        # Count of elements with coordinate-compressed index >= threshold_idx
        count = bit_range_sum(threshold_idx, num_values - 1)
        ans += count
        
        # Add S[r+1] (can be used as l for the next r+1)
        bit_add(value_to_idx[S[r + 1]])
    
    print(ans)

solve()

---

This editorial was generated by claude4.5opus.