O(NlogN) Python solution, binary indexed tree, 268 ms


  • 6
    K

    Sum[k] is the sum of first k numbers. O(N^2) solution is

    for j in range(n + 1):
        for i in range(j):
            if lower <= Sum[j] - Sum[i] <= upper: res += 1
    

    This is equal to:

    collection = empty
    for sum_j in Sum:
        sum_i_count = how many sum_i in this collection that sum_j - upper <= sum_i <= sum_j - lower
        res += sum_i_count
        put sum_j into this collection
    

    With Binary indexed tree, counting sum_i number is O(logN), putting sum_i into tree is also O(logN). Here we store the index of sortSum in the tree. Since index of BITree starts from 1, we need bisect.bisect_left(sortSum, sum_j) + 1 for update().

    def countRangeSum(self, nums, lower, upper):
        n = len(nums)
        Sum, BITree = [0] * (n + 1), [0] * (n + 2)
        
        def count(x):
            s = 0
            while x:
                s += BITree[x]
                x -= (x & -x)
            return s
        
        def update(x):
            while x <= n + 1:
                BITree[x] += 1
                x += (x & -x)
                
        for i in range(n):
            Sum[i + 1] = Sum[i] + nums[i]
        sortSum, res = sorted(Sum), 0
        for sum_j in Sum:
            sum_i_count = count(bisect.bisect_right(sortSum, sum_j - lower)) - count(bisect.bisect_left(sortSum, sum_j - upper))
            res += sum_i_count
            update(bisect.bisect_left(sortSum, sum_j) + 1)
        return res
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.