Count of Range Sum — Hard Problem & Solution

The range sum S(i, j) is the sum of nums[i..j] for i <= j. Given nums and two integers lower and upper, return the number of index pairs (i, j) for which…

Difficulty: Hard
Topics: Prefix Sum, Divide and Conquer, Binary Indexed Tree
Asked at: Amazon, Google, Meta
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

The range sum S(i, j) is the sum of nums[i..j] for i <= j.

Given nums and two integers lower and upper, return the number of index pairs (i, j) for which S(i, j) lies in [lower, upper] inclusive.

Example 1

Input: nums = [-2,5,-1], lower = -2, upper = 2
Output: 3
Explanation: The qualifying ranges are [0,0], [2,2] and [0,2], summing to -2, -1 and 2.

Example 2

Input: nums = [0], lower = 0, upper = 0
Output: 1

Example 3

Input: nums = [1,2,3], lower = 10, upper = 20
Output: 0

Constraints

1 <= nums.length <= 10000
-10000 <= nums[i] <= 10000
-100000000 <= lower <= upper <= 100000000

How to solve Count of Range Sum

Turn range sums into prefix-sum differences, then sweep left to right asking a counting question about the prefixes already seen. Coordinate compression puts the (possibly huge) prefix values into a Fenwick tree's index space.

Approach

Build the prefix array P of length n + 1 with P[0] = 0.
Sort and deduplicate P to get the compression table uniq.
Insert P[0] into a Fenwick tree over uniq's ranks.
For j from 1 to n: count the inserted prefixes in [P[j] - upper, P[j] - lower] using two binary searches into uniq and two Fenwick prefix queries, then insert P[j].

Why it works

P[j] - P[i] lies in [lower, upper] exactly when P[i] lies in [P[j] - upper, P[j] - lower]. Processing j in increasing order means the tree holds precisely the valid left endpoints i < j, so each qualifying pair is counted once at its right end.

Complexity

Time — O(n log n)
Space — O(n)

Pitfalls

The query bounds P[j] - upper and P[j] - lower may not appear in uniq, so the binary searches must be a strict lower bound and a non-strict upper bound rather than exact lookups.
A Fenwick tree is 1-indexed; adding 1 to the rank is mandatory or the update loop never terminates.
Prefix sums and the shifted bounds exceed 32 bits once n and the values grow — accumulate in 64 bits.
Merge sort over the prefix array is the other standard solution with the same complexity.

Reference solution

Python

from bisect import bisect_left, bisect_right
from typing import List

def countRangeSum(nums: List[int], lower: int, upper: int) -> int:
    n = len(nums)
    P = [0] * (n + 1)
    for i in range(n):
        P[i + 1] = P[i] + nums[i]
    uniq = sorted(set(P))
    m = len(uniq)
    bit = [0] * (m + 1)

    def add(pos: int) -> None:
        while pos <= m:
            bit[pos] += 1
            pos += pos & -pos

    def pref(pos: int) -> int:
        s = 0
        while pos > 0:
            s += bit[pos]
            pos -= pos & -pos
        return s

    total = 0
    add(bisect_left(uniq, P[0]) + 1)
    for j in range(1, n + 1):
        total += pref(bisect_right(uniq, P[j] - lower)) - pref(bisect_left(uniq, P[j] - upper))
        add(bisect_left(uniq, P[j]) + 1)
    return total

JavaScript

var countRangeSum = function(nums, lower, upper) {
    var n = nums.length;
    var P = [0];
    for (var i = 0; i < n; i++) P.push(P[i] + nums[i]);
    var sorted = P.slice().sort(function(a, b) { return a - b; });
    var uniq = [];
    for (var t = 0; t < sorted.length; t++) {
        if (t === 0 || sorted[t] !== sorted[t - 1]) uniq.push(sorted[t]);
    }
    var m = uniq.length;
    var bit = [];
    for (var u = 0; u <= m; u++) bit.push(0);
    var add = function(pos) { for (var p = pos; p <= m; p += p & -p) bit[p]++; };
    var pref = function(pos) { var s = 0; for (var p = pos; p > 0; p -= p & -p) s += bit[p]; return s; };
    var lowerBound = function(v) {
        var lo = 0, hi = m;
        while (lo < hi) { var mid = (lo + hi) >> 1; if (uniq[mid] < v) lo = mid + 1; else hi = mid; }
        return lo;
    };
    var upperBound = function(v) {
        var lo = 0, hi = m;
        while (lo < hi) { var mid = (lo + hi) >> 1; if (uniq[mid] <= v) lo = mid + 1; else hi = mid; }
        return lo;
    };
    var total = 0;
    add(lowerBound(P[0]) + 1);
    for (var j = 1; j <= n; j++) {
        total += pref(upperBound(P[j] - lower)) - pref(lowerBound(P[j] - upper));
        add(lowerBound(P[j]) + 1);
    }
    return total;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 63 prefix sum problems · the whole catalogue