Problem statement

https://leetcode.com/problems/count-of-range-sum/

Solution 1

There is very short solutoin using BST. Let us first evaluate all cumulative sums S0, S1, ... S_{n-1}. Then the algorithm is to add sums one by one to BST (SortedList) and for each new element, before we add it, do binary search in our SortedList of two borders.

Complexity

Overall time complexity is O(n log n), space complexity is O(n).

Code

from sortedcontainers import SortedList

class Solution:
    def countRangeSum(self, nums, lower, upper):
        SList, ans = SortedList([0]), 0

        for s in accumulate(nums):
            ans += SList.bisect_right(s - lower) - SList.bisect_left(s - upper)
            SList.add(s)
        return ans

Solution 2

Here is solution using BIT, see problem 493 for more details. The idea is to create BIT with n + 1 elements, for example for cumulative sums [2, 4, 4, 3, 6, 1] when we add element by element we will have find the rank of element, so elements in our BIT will look like: [0, 0, 1, 0, 0, 0, 0] -> [0, 0, 1, 0, 0, 1, 0] -> [0, 0, 1, 0, 0, 2, 0] -> [0, 0, 1, 1, 0, 2, 0] -> [0, 0, 1, 1, 0, 2, 1] -> [0, 1, 1, 1, 0, 2, 1] (it is not real elements, but what list this BIT represents). Then to find how many good elements we have, we need to evaluate sum in range: we need to find range using binary search, for example for number 5 we need to find where in our BIT we should end: it is just after element 4 in our cumulative sums.

Complexity

It is also O(n log n) for time and O(n) for space.

Code

from bisect import bisect, bisect_left

class BIT:
    def __init__(self, n):
        self.sums = [0] * (n+1)
    
    def update(self, i, delta):
        while i < len(self.sums):
            self.sums[i] += delta
            i += i & (-i)
    
    def query(self, i):
        res = 0
        while i > 0:
            res += self.sums[i]
            i -= i & (-i)
        return res

class Solution:
    def countRangeSum(self, nums, lower, upper):
        n = len(nums)
        acc = [0] + list(accumulate(nums))
        bit = BIT(n + 1)
        acc_sorted = sorted(acc)
        
        res = 0
        for s in acc:
            left = bit.query(bisect_left(acc_sorted, s - upper))
            right = bit.query(bisect(acc_sorted, s - lower))
            res += right - left
            bit.update(bisect(acc_sorted, s), 1)
        return res

Remark

There is Divide and Conquer approach with O(n log^2n) complexity: we divide numbers into two parts, solve problem for left and right sub-problems. Now we need to solve it in the middle: range in the middle will be sum of two ranges. We can evaluate all possible values for left and right range in O(n), sort them in O(n log n) and then evaluate number of combined ranges in O(n log n) also. Recurrent equation is T(n) = 2T(n/2) + O(n log n) which gives us O(n log^2n) complexity.

Similar problems:

0315 Count of Smaller Number After Self.

0493 Reverse Pairs