Count Subarrays with Sum in a Range
Given an integer array $A_1, A_2, \dots, A_n$ (which may contain negative numbers) and integer bounds $L$ and $R$, count the number of contiguous subarrays whose sum falls in the range $[L, R]$.
Design an $O(n \log n)$ algorithm. Your solution should handle negative numbers, 64-bit overflow, and use coordinate compression for stability.
Constraints: - $1 \leq n \leq 10^5$ - $-10^9 \leq A_i \leq 10^9$ - $-10^{18} \leq L \leq R \leq 10^{18}$
Example 1: - Input: $A = [1, 2, -1, 3]$, $L = 2$, $R = 4$ - Output: $6$ - Explanation: The subarrays with sums in $[2, 4]$ are: $[1,2]$ (sum 3), $[1,2,-1]$ (sum 2), $[2]$ (sum 2), $[2,-1,3]$ (sum 4), $[-1,3]$ (sum 2), $[3]$ (sum 3). That gives 6 valid subarrays.
Example 2: - Input: $A = [-1, -1, 1]$, $L = -1$, $R = 0$ - Output: $4$ - Explanation: The subarrays with sums in $[-1, 0]$ are: $[-1]$ (index 0, sum -1), $[-1]$ (index 1, sum -1), $[-1, 1]$ (sum 0), and $[-1, -1, 1]$ (sum -1). That gives 4 valid subarrays.
Hints
- Rewrite the subarray sum $\sum_{k=i+1}^{j} A_k$ as a difference of prefix sums $P_j - P_i$. The problem becomes: count pairs $(i, j)$ with $i < j$ and $P_j - P_i \in [L, R]$.
- For each $j$, you need the count of previously seen prefix sums in the range $[P_j - R, P_j - L]$. This is a range-counting query -- use a Fenwick tree (BIT) or merge sort.
- Coordinate-compress the prefix sum values to keep the Fenwick tree size $O(n)$. Use binary search on the sorted values to map query bounds to ranks.
Worked Solution
How to Think About It: The brute force approach checks all $O(n^2)$ subarrays, which is too slow for $n = 10^5$. The key observation is that the sum of subarray $A[i+1..j]$ equals $P_j - P_i$ where $P$ is the prefix sum array. So counting subarrays with sum in $[L, R]$ reduces to counting pairs $(i, j)$ with $i < j$ and $L \leq P_j - P_i \leq R$, i.e., $P_j - R \leq P_i \leq P_j - L$. For each $j$, you count how many earlier prefix sums fall in a specific range -- a range-counting query.
Algorithm:
1. Compute prefix sums $P_0 = 0,\, P_j = P_{j-1} + A_j$ for $j = 1, \dots, n$. Use 64-bit integers to avoid overflow. 2. Coordinate-compress the prefix sum values: collect all $P_i$ values, sort and deduplicate them, and map each to a rank. 3. Initialize a Fenwick tree (BIT) over the compressed ranks. 4. Process prefix sums left to right. For each $j = 0, 1, \dots, n$: - Query the Fenwick tree for the count of previously inserted prefix sums in $[P_j - R,\, P_j - L]$, using binary search on the sorted compressed values to find the rank range. - Add $P_j$ to the Fenwick tree at its compressed rank. 5. The total query count is the answer.
Code:
```python import bisect
def count_subarrays_in_range(A, L, R): n = len(A) prefix = [0] * (n + 1) for i in range(n): prefix[i + 1] = prefix[i] + A[i] sorted_vals = sorted(set(prefix)) size = len(sorted_vals) rank = {v: i + 1 for i, v in enumerate(sorted_vals)} tree = [0] * (size + 1) def update(i): while i <= size: tree[i] += 1 i += i & (-i) def query(i): s = 0 while i > 0: s += tree[i] i -= i & (-i) return s count = 0 for j in range(n + 1): lo = prefix[j] - R hi = prefix[j] - L lr = bisect.bisect_left(sorted_vals, lo) hr = bisect.bisect_right(sorted_vals, hi) if lr < hr: count += query(hr) - query(lr) update(rank[prefix[j]]) return count ```
Handling edge cases:
- *64-bit overflow:* Prefix sums can reach $n \times 10^9 = 10^{14}$, which fits in a 64-bit integer. In C++/Java, use
long long/long. - *Negative numbers:* The algorithm works unchanged because prefix sums are not monotone; coordinate compression handles arbitrary orderings.
- *Coordinate compression stability:* Compressing only the actual prefix sum values (not the query bounds) avoids issues with bounds that fall outside the observed range; the binary search handles this gracefully.
Complexity: - Time: $O(n \log n)$ -- one Fenwick update and one range query per prefix sum, each $O(\log n)$. - Space: $O(n)$.
Answer: Use prefix sums to reduce the problem to counting pairs with differences in $[L, R]$, then sweep left-to-right with a Fenwick tree over coordinate-compressed ranks for $O(n \log n)$ total time.
Intuition
The prefix sum transformation is the standard trick for turning subarray sum problems into pair-counting problems. Once you see that $\text{sum}(A[i+1..j]) = P_j - P_i$, the problem reduces to: as you scan $j$ from left to right, how many of the previously seen $P_i$ values fall in the window $[P_j - R, P_j - L]$? This is exactly a dynamic range-counting query, which a Fenwick tree handles in $O(\log n)$ per operation.
The subtlety with negative numbers is that prefix sums are not monotonically increasing, so you cannot use a two-pointer approach. Coordinate compression solves the problem of having prefix sums spread across a huge range (up to $10^{14}$) -- by mapping them to consecutive ranks, the Fenwick tree stays $O(n)$ in size. An alternative $O(n \log n)$ approach uses merge sort to count inversions in a modified sense, but the Fenwick tree approach is more intuitive and easier to implement correctly.