LeetCode 1508: Range Sum of Sorted Subarray Sums

Brute Force

We can generate all $\mathcal{O}(n^2)$ sums in $\mathcal{O}(n^2)$ time. Sort them in $\mathcal{O}( n^2 \cdot \lg{n} )$ time. Then, return (right-left+1) of them in [left : right+1].

Time would be $\mathcal{O}( n^2 \cdot \lg{n} )$ .

Generate sums in sorted order

If we could generate the sub-array sums one at a time in sorted order, we could generate first (left-1) of them and throw away. The next (right-left+1) sums would be our answer.

With only positive numbers, the sequence of subarray sums starting at a particular index can be treated as a sorted linked list. We put the head of each possible starting index in a min heap and pop the top and (if any) push the next from that list.

Time: $\mathcal{O}( \max{ \left( n, \texttt{right} \cdot \lg{n} \right) } )$ . Space: $\mathcal{O}( \max{ (n, \texttt{right} - \texttt{left} + 1) } )$ .

from heapq import heapify, heappush as push, heappop as pop

class Solution:
    
    def sum_top_k(self, minq, k, nums, MOD) -> int:
        n = len(nums)
        ksum = 0
        while k > 0:
            k -= 1
            s, i = pop( minq )
            ksum = (ksum + s) % MOD
            if i == n-1:
                continue
            i += 1
            s = (s + nums[i]) % MOD
            push( minq, (s, i) )

        return ksum
            

    def rangeSum(self, nums: List[int], n: int, left: int, right: int) -> int:
        MOD = pow(10, 9)+7

        minq = [
            (x, i) for i, x in enumerate(nums)
        ]
        heapify(minq)

        _ = self.sum_top_k( minq, left-1, nums, MOD )
        
        return self.sum_top_k( minq, right-left+1, nums, MOD )

unreasonably effective