Problem statement

https://leetcode.com/problems/range-sum-of-sorted-subarray-sums/

Solution

You can just calculate sums, sort them and find in range here, constraints allow to do it.

However we can do better. Notice that if we evaluate cumulative sums A = acc, then if we create B = -acc, then what we have is sums of pairs of elements, one from A and another from B. Now, idea is to use BinarySearch 0379 K Largest Pairs problem. Let check(k) be the sum of k-th largest numbers. We can use binary search to find such sum: first, we find value beg, such that sums[beg] = k-th largest number. Then we need to check how many times we have repeating value and adjust sum. Finally, we need to evaluate check(n(n-1)//2 - i + 1) - check(n(n-1)//2 - j), because we have n*(n-1)` zeroes and negative numbers and we interested only in positive.

Complexity

It is O(n log M) for time, where M = sum(nums) and O(n) for time.

Code

class Solution:
    def rangeSum(self, nums, n, i, j):
        acc = [0] + list(accumulate(nums))
        A = acc[:]
        B = [-x for x in acc[::-1]]
        acc_B = [0] + list(accumulate(B))
        n = len(A)
        
        def check(k):
            end, beg = A[-1] + B[-1], A[0] + B[0]
            while beg < end:
                mid = (beg + end) // 2
                j, cnt = len(B) - 1, 0
                for i in range(len(A)):
                    while j >= 0 and A[i] + B[j] > mid:
                        j -= 1
                    cnt += j + 1
                
                if cnt < k:
                    beg = mid + 1
                else:
                    end = mid

            ans, taken = 0, 0
            
            for x in A:
                idx = bisect.bisect(B, beg - x)
                taken += idx
                ans += acc_B[idx] + x*idx
            return -ans + (taken - k)*beg

        return (check(n*(n-1)//2 - i + 1) - check(n*(n-1)//2 - j)) % (10**9 + 7)