Problem statement

https://binarysearch.com/problems/Sum-of-Sublist-Range-Sum/

Solution

Equal to Leetcode 1508. Range Sum of Sorted Subarray Sums.

Complexity

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

Code

class Solution:
    def solve(self, nums, 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) - check(n*(n-1)//2 - j - 1)