Count of Range Sum

Hard

#0314Time: O(n^2). The two nested loops run `n * (n+1) / 2` times, leading to a quadratic time complexity. This will be too slow for `n = 10^5`.Space: O(n). We need an additional array of size `n+1` to store the prefix sums.1 company

Algorithms

Binary Search Divide and Conquer Merge Sort

Data structures

Array Binary Indexed Tree Segment Tree Ordered Set

Companies

Morgan Stanley

Prompt

Given an integer array nums and two integers lower and upper, return the number of range sums that lie in [lower, upper] inclusive.

Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j inclusive, where i <= j.

Example 1:

Input: nums = [-2,5,-1], lower = -2, upper = 2
Output: 3
Explanation: The three ranges are: [0,0], [2,2], and [0,2] and their respective sums are: -2, -1, 2.

Example 2:

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

Constraints:

1 <= nums.length <= 10⁵
-2³¹ <= nums[i] <= 2³¹ - 1
-10⁵ <= lower <= upper <= 10⁵
The answer is guaranteed to fit in a 32-bit integer.

Approaches

3 approaches with complexity analysis and trade-offs.

This approach is a straightforward improvement over the naive O(n^3) method. The core idea is to first compute the prefix sums of the input array. The sum of any range [i, j] can then be calculated in O(1) time as prefixSum[j+1] - prefixSum[i]. We can then iterate through all possible start and end indices (i, j) and check if their range sum falls within [lower, upper].

Algorithm

Create a long array prefix of size n + 1 and initialize prefix[0] = 0.
Iterate from i = 0 to n-1, calculating prefix[i+1] = prefix[i] + nums[i].
Initialize count = 0.
Use a nested loop. The outer loop for i from 0 to n-1.
The inner loop for j from i to n-1.
Calculate sum = prefix[j+1] - prefix[i].
If sum >= lower && sum <= upper, increment count.
Return count.

Walkthrough

First, we create a prefix sum array, let's call it prefix, of size n+1. We use long for the prefix sums to prevent potential integer overflow, as the sum of up to 10^5 numbers can exceed the int range. prefix[0] is initialized to 0. We populate this array such that prefix[k] stores the sum of the first k elements of nums. So, prefix[k] = nums[0] + ... + nums[k-1]. Then, we use two nested loops to iterate through all possible subarrays. The outer loop iterates from i = 0 to n-1 (the start of the range), and the inner loop iterates from j = i to n-1 (the end of the range). For each pair (i, j), we calculate the range sum S(i, j) using prefix[j+1] - prefix[i]. We check if this sum is between lower and upper, inclusive. If it is, we increment a counter. After checking all pairs, the counter holds the total number of valid range sums.

1public int countRangeSum(int[] nums, int lower, int upper) {2    int n = nums.length;3    long[] prefix = new long[n + 1];4    for (int i = 0; i < n; i++) {5        prefix[i + 1] = prefix[i] + nums[i];6    }7 8    int count = 0;9    for (int i = 0; i < n; i++) {10        for (int j = i; j < n; j++) {11            long sum = prefix[j + 1] - prefix[i];12            if (sum >= lower && sum <= upper) {13                count++;14            }15        }16    }17    return count;18}

Complexity

Time

O(n^2). The two nested loops run `n * (n+1) / 2` times, leading to a quadratic time complexity. This will be too slow for `n = 10^5`.

Space

O(n). We need an additional array of size `n+1` to store the prefix sums.

Trade-offs

Pros

Simple to understand and implement.
Improves upon the O(n^3) naive solution.

Cons

Inefficient for large inputs and will result in a "Time Limit Exceeded" error for the given constraints.

This approach also achieves O(n log n) complexity. After transforming the problem using prefix sums (prefix[j] - upper <= prefix[i] <= prefix[j] - lower), we can iterate through the prefix sums. For each prefix[j], we need to count how many previous sums prefix[i] fall into the required range. A Binary Indexed Tree (BIT) or Fenwick Tree is well-suited for this. Since the values of prefix sums can be large and non-contiguous, we first need to discretize (or perform coordinate compression on) all relevant values.

Algorithm

Calculate the long prefix sum array prefix.
Create a set of all relevant values: p, p-lower, p-upper for each p in prefix.
Sort the unique values from the set to create a ranked list. Build a map from value to its rank.
Initialize a Fenwick Tree (BIT) of size equal to the number of unique values.
Initialize count = 0.
Iterate through each sum p in the prefix array.
For the current p, find the ranks for the range boundaries p-upper and p-lower.
Query the BIT to get the number of elements in the valid range and add it to count.
Update the BIT at the rank of p: bit.update(rank(p), 1).
Return count.

Walkthrough

1. Prefix Sums: Calculate the prefix sum array prefix of size n+1. 2. Discretization: The values in prefix can be large, so we can't use them as indices for the BIT. We collect all values that we will need to handle: prefix[i], prefix[i] - lower, and prefix[i] - upper for all i. We gather all these values, find the unique ones, sort them, and create a map from each unique value to its rank (index in the sorted list). This allows us to work with small, contiguous indices. 3. BIT Operations: We initialize a BIT with a size equal to the number of unique values. We iterate through the prefix sums from j = 0 to n. For each prefix[j]:

First, we query the BIT to find the count of previous sums. We need to count prefix[i] (i < j) in the range [prefix[j] - upper, prefix[j] - lower]. We find the ranks of the range boundaries, say rank_low and rank_high. The count is then BIT.query(rank_high) - BIT.query(rank_low - 1). We add this to our total count.
Second, we update the BIT with the current prefix sum prefix[j]. We find its rank, rank_current, and call BIT.update(rank_current, 1).

1// Note: This is a conceptual representation. A full implementation2// requires a BIT class and careful handling of discretization.3public int countRangeSum(int[] nums, int lower, int upper) {4    int n = nums.length;5    long[] prefix = new long[n + 1];6    for (int i = 0; i < n; i++) {7        prefix[i + 1] = prefix[i] + nums[i];8    }9 10    // Discretization11    Set<Long> valueSet = new HashSet<>();12    for (long p : prefix) {13        valueSet.add(p);14        valueSet.add(p - lower);15        valueSet.add(p - upper);16    }17    List<Long> sortedValues = new ArrayList<>(valueSet);18    Collections.sort(sortedValues);19    Map<Long, Integer> ranks = new HashMap<>();20    for (int i = 0; i < sortedValues.size(); i++) {21        ranks.put(sortedValues.get(i), i + 1); // 1-based index for BIT22    }23 24    // BIT25    FenwickTree bit = new FenwickTree(ranks.size());26    int count = 0;27    for (long p : prefix) {28        // Count previous sums in [p - upper, p - lower]29        int rankHigh = ranks.get(p - lower);30        int rankLowBound = ranks.get(p - upper);31        count += bit.query(rankHigh) - bit.query(rankLowBound - 1);32        // Add current sum to BIT33        bit.update(ranks.get(p), 1);34    }35    return count;36}37 38// FenwickTree (BIT) helper class must be implemented separately39class FenwickTree {40    private int[] tree;41    private int size;42    public FenwickTree(int size) { /* ... */ }43    public void update(int index, int delta) { /* ... */ }44    public int query(int index) { /* ... */ }45}

Complexity

Time

O(n log n). Discretization involves sorting, which is `O(n log n)`. The main loop runs `n+1` times, with each BIT operation taking `O(log k)` where `k` is the number of unique values (`k <= 3(n+1)`). So the loop is also `O(n log n)`.

Space

O(n). For prefix sums, the set of values, the rank map, and the BIT.

Trade-offs

Pros

Provides an alternative O(n log n) solution.
The iterative structure might be more intuitive for some than recursion.

Cons

Implementation is more complex due to the need for discretization and a BIT data structure.
The constant factors might be higher than the merge sort approach due to overhead from data structures (HashMap, BIT).

This is a highly efficient and canonical approach that solves the problem in O(n log n) time. The problem is first transformed using prefix sums. We need to find the number of pairs (i, j) with i < j such that lower <= prefix[j] - prefix[i] <= upper. This is equivalent to finding pairs where prefix[j] - upper <= prefix[i] <= prefix[j] - lower. This condition can be solved efficiently during the merge step of a merge sort algorithm applied to the prefix sum array.

Algorithm

Create a long prefix sum array prefix of size n+1.
Define a recursive function countWhileMerging(sums, start, end, lower, upper).
Base Case: If end <= start, return 0.
Divide: Find mid = start + (end - start) / 2.
Conquer: count = countWhileMerging(..., start, mid, ...) + countWhileMerging(..., mid + 1, end, ...).
Count cross-pairs: For each i from start to mid, find the number of j's from mid+1 to end such that sums[j] is in [sums[i] + lower, sums[i] + upper]. Use two pointers on the sorted right half to do this in O(end - start) total time. Add this to count.
Merge: Merge the sorted halves sums[start...mid] and sums[mid+1...end] into a single sorted array.
Return count.

Walkthrough

First, we compute the prefix sum array prefix of size n+1. The core of the solution is a recursive function, say countWhileMerging, that takes a subarray of the prefix sums prefix[start...end].

Divide: The function splits the array into two halves: [start...mid] and [mid+1...end].
Conquer: It recursively calls itself on both halves. The total count is the sum of counts from the left half, the right half, and the counts of pairs (i, j) where i is in the left half and j is in the right half.
Combine (Count & Merge): This is the key step. For each element prefix[j] in the right half, we need to count how many elements prefix[i] in the left half satisfy prefix[j] - upper <= prefix[i] <= prefix[j] - lower. Since the recursive calls ensure that the left and right halves are sorted independently, we can find these prefix[i]s efficiently. We use two pointers, k and l, on the left half to find the range of valid prefix[i]s for each prefix[j]. As j iterates through the right half, k and l only move forward, making this counting step linear in time, i.e., O(end - start).
After counting, we perform a standard merge operation to sort the prefix[start...end] subarray, which also takes linear time. The recurrence relation T(n) = 2T(n/2) + O(n) solves to O(n log n).

1public int countRangeSum(int[] nums, int lower, int upper) {2    int n = nums.length;3    long[] prefixSums = new long[n + 1];4    for (int i = 0; i < n; i++) {5        prefixSums[i + 1] = prefixSums[i] + nums[i];6    }7    return countWhileMerging(prefixSums, 0, n, lower, upper);8}9 10private int countWhileMerging(long[] sums, int start, int end, int lower, int upper) {11    if (end <= start) {12        return 0;13    }14    int mid = start + (end - start) / 2;15    int count = countWhileMerging(sums, start, mid, lower, upper)16              + countWhileMerging(sums, mid + 1, end, lower, upper);17 18    // Count pairs (i, j) with i in left half, j in right half19    int k = mid + 1, l = mid + 1;20    for (int i = start; i <= mid; i++) {21        // Find first k s.t. sums[k] >= sums[i] + lower22        while (k <= end && sums[k] < sums[i] + lower) {23            k++;24        }25        // Find first l s.t. sums[l] > sums[i] + upper26        while (l <= end && sums[l] <= sums[i] + upper) {27            l++;28        }29        count += l - k;30    }31 32    // Merge step33    long[] merged = new long[end - start + 1];34    int p1 = start, p2 = mid + 1, p = 0;35    while (p1 <= mid || p2 <= end) {36        if (p1 <= mid && (p2 > end || sums[p1] <= sums[p2])) {37            merged[p++] = sums[p1++];38        } else {39            merged[p++] = sums[p2++];40        }41    }42    System.arraycopy(merged, 0, sums, start, merged.length);43 44    return count;45}

Complexity

Time

O(n log n). The algorithm follows the merge sort recurrence.

Space

O(n). This is required for the temporary array used during the merge step.

Trade-offs

Pros

Very efficient and passes for large constraints.
Elegant divide-and-conquer solution.
Doesn't require complex data structures like Fenwick or Segment trees.

Cons

Can be tricky to implement correctly, especially the counting logic within the merge step.
The recursion adds some overhead.

Solutions

Solution

1class BinaryIndexedTree { private int n ; private int [] c ; public BinaryIndexedTree ( int n ) { this . n = n ; this . c = new int [ n + 1 ]; } public void update ( int x , int v ) { while ( x <= n ) { c [ x ] += v ; x += x & - x ; } } public int query ( int x ) { int s = 0 ; while ( x != 0 ) { s += c [ x ]; x -= x & - x ; } return s ; } } class Solution { public int countRangeSum ( int [] nums , int lower , int upper ) { int n = nums . length ; long [] s = new long [ n + 1 ]; for ( int i = 0 ; i < n ; ++ i ) { s [ i + 1 ] = s [ i ] + nums [ i ]; } long [] arr = new long [ n * 3 + 3 ]; for ( int i = 0 , j = 0 ; i <= n ; ++ i , j += 3 ) { arr [ j ] = s [ i ]; arr [ j + 1 ] = s [ i ] - lower ; arr [ j + 2 ] = s [ i ] - upper ; } Arrays . sort ( arr ); int m = 0 ; for ( int i = 0 ; i < arr . length ; ++ i ) { if ( i == 0 || arr [ i ] != arr [ i - 1 ]) { arr [ m ++] = arr [ i ]; } } BinaryIndexedTree tree = new BinaryIndexedTree ( m ); int ans = 0 ; for ( long x : s ) { int l = search ( arr , m , x - upper ); int r = search ( arr , m , x - lower ); ans += tree . query ( r ) - tree . query ( l - 1 ); tree . update ( search ( arr , m , x ), 1 ); } return ans ; } private int search ( long [] nums , int r , long x ) { int l = 0 ; while ( l < r ) { int mid = ( l + r ) >> 1 ; if ( nums [ mid ] >= x ) { r = mid ; } else { l = mid + 1 ; } } return l + 1 ; } }

1class BinaryIndexedTree { public: BinaryIndexedTree ( int _n ) : n ( _n ) , c ( _n + 1 ) {} void update ( int x , int v ) { while ( x <= n ) { c [ x ] += v ; x += x & - x ; } } int query ( int x ) { int s = 0 ; while ( x ) { s += c [ x ]; x -= x & - x ; } return s ; } private: int n ; vector < int > c ; }; class Solution { public: int countRangeSum ( vector < int >& nums , int lower , int upper ) { using ll = long long ; int n = nums . size (); ll s [ n + 1 ]; s [ 0 ] = 0 ; for ( int i = 0 ; i < n ; ++ i ) { s [ i + 1 ] = s [ i ] + nums [ i ]; } ll arr [( n + 1 ) * 3 ]; for ( int i = 0 , j = 0 ; i <= n ; ++ i , j += 3 ) { arr [ j ] = s [ i ]; arr [ j + 1 ] = s [ i ] - lower ; arr [ j + 2 ] = s [ i ] - upper ; } sort ( arr , arr + ( n + 1 ) * 3 ); int m = unique ( arr , arr + ( n + 1 ) * 3 ) - arr ; BinaryIndexedTree tree ( m ); int ans = 0 ; for ( int i = 0 ; i <= n ; ++ i ) { int l = lower_bound ( arr , arr + m , s [ i ] - upper ) - arr + 1 ; int r = lower_bound ( arr , arr + m , s [ i ] - lower ) - arr + 1 ; ans += tree . query ( r ) - tree . query ( l - 1 ); tree . update ( lower_bound ( arr , arr + m , s [ i ]) - arr + 1 , 1 ); } return ans ; } };

1class BinaryIndexedTree : def __init__ ( self , n ): self . n = n self . c = [ 0 ] * ( n + 1 ) def update ( self , x , v ): while x <= self . n : self . c [ x ] += v x += x & - x def query ( self , x ): s = 0 while x > 0 : s += self . c [ x ] x -= x & - x return s class Solution : def countRangeSum ( self , nums : List [ int ], lower : int , upper : int ) -> int : s = list ( accumulate ( nums , initial = 0 )) arr = sorted ( set ( v for x in s for v in ( x , x - lower , x - upper ))) tree = BinaryIndexedTree ( len ( arr )) ans = 0 for x in s : l = bisect_left ( arr , x - upper ) + 1 r = bisect_left ( arr , x - lower ) + 1 ans += tree . query ( r ) - tree . query ( l - 1 ) tree . update ( bisect_left ( arr , x ) + 1 , 1 ) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem

Solution

1class BinaryIndexedTree { private int n ; private int [] c ; public BinaryIndexedTree ( int n ) { this . n = n ; this . c = new int [ n + 1 ]; } public void update ( int x , int v ) { while ( x <= n ) { c [ x ] += v ; x += x & - x ; } } public int query ( int x ) { int s = 0 ; while ( x != 0 ) { s += c [ x ]; x -= x & - x ; } return s ; } } class Solution { public int countRangeSum ( int [] nums , int lower , int upper ) { int n = nums . length ; long [] s = new long [ n + 1 ]; for ( int i = 0 ; i < n ; ++ i ) { s [ i + 1 ] = s [ i ] + nums [ i ]; } long [] arr = new long [ n * 3 + 3 ]; for ( int i = 0 , j = 0 ; i <= n ; ++ i , j += 3 ) { arr [ j ] = s [ i ]; arr [ j + 1 ] = s [ i ] - lower ; arr [ j + 2 ] = s [ i ] - upper ; } Arrays . sort ( arr ); int m = 0 ; for ( int i = 0 ; i < arr . length ; ++ i ) { if ( i == 0 || arr [ i ] != arr [ i - 1 ]) { arr [ m ++] = arr [ i ]; } } BinaryIndexedTree tree = new BinaryIndexedTree ( m ); int ans = 0 ; for ( long x : s ) { int l = search ( arr , m , x - upper ); int r = search ( arr , m , x - lower ); ans += tree . query ( r ) - tree . query ( l - 1 ); tree . update ( search ( arr , m , x ), 1 ); } return ans ; } private int search ( long [] nums , int r , long x ) { int l = 0 ; while ( l < r ) { int mid = ( l + r ) >> 1 ; if ( nums [ mid ] >= x ) { r = mid ; } else { l = mid + 1 ; } } return l + 1 ; } }

1class BinaryIndexedTree { public: BinaryIndexedTree ( int _n ) : n ( _n ) , c ( _n + 1 ) {} void update ( int x , int v ) { while ( x <= n ) { c [ x ] += v ; x += x & - x ; } } int query ( int x ) { int s = 0 ; while ( x ) { s += c [ x ]; x -= x & - x ; } return s ; } private: int n ; vector < int > c ; }; class Solution { public: int countRangeSum ( vector < int >& nums , int lower , int upper ) { using ll = long long ; int n = nums . size (); ll s [ n + 1 ]; s [ 0 ] = 0 ; for ( int i = 0 ; i < n ; ++ i ) { s [ i + 1 ] = s [ i ] + nums [ i ]; } ll arr [( n + 1 ) * 3 ]; for ( int i = 0 , j = 0 ; i <= n ; ++ i , j += 3 ) { arr [ j ] = s [ i ]; arr [ j + 1 ] = s [ i ] - lower ; arr [ j + 2 ] = s [ i ] - upper ; } sort ( arr , arr + ( n + 1 ) * 3 ); int m = unique ( arr , arr + ( n + 1 ) * 3 ) - arr ; BinaryIndexedTree tree ( m ); int ans = 0 ; for ( int i = 0 ; i <= n ; ++ i ) { int l = lower_bound ( arr , arr + m , s [ i ] - upper ) - arr + 1 ; int r = lower_bound ( arr , arr + m , s [ i ] - lower ) - arr + 1 ; ans += tree . query ( r ) - tree . query ( l - 1 ); tree . update ( lower_bound ( arr , arr + m , s [ i ]) - arr + 1 , 1 ); } return ans ; } };

1class BinaryIndexedTree : def __init__ ( self , n ): self . n = n self . c = [ 0 ] * ( n + 1 ) def update ( self , x , v ): while x <= self . n : self . c [ x ] += v x += x & - x def query ( self , x ): s = 0 while x > 0 : s += self . c [ x ] x -= x & - x return s class Solution : def countRangeSum ( self , nums : List [ int ], lower : int , upper : int ) -> int : s = list ( accumulate ( nums , initial = 0 )) arr = sorted ( set ( v for x in s for v in ( x , x - lower , x - upper ))) tree = BinaryIndexedTree ( len ( arr )) ans = 0 for x in s : l = bisect_left ( arr , x - upper ) + 1 r = bisect_left ( arr , x - lower ) + 1 ans += tree . query ( r ) - tree . query ( l - 1 ) tree . update ( bisect_left ( arr , x ) + 1 , 1 ) return ans

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.