Problem statement

https://leetcode.com/problems/kth-smallest-subarray-sum/

Solution

The idea is to given number T to have function check(T) which answer the question: how many sums of subarrays <= T. To answer this question we need to calculate cumulative sums of numbers and then use two pointers approach, similar to 2sum problem.

Next step is to use binary search for T from 0 to sum(nums): we know that answer check(0) = False, because all numbers are positive and there no sums of subarrays which <= 0. Also we know that answer of check(sum(num)) = True, because every sum of subarray less then sum of all elements. So, we have pattern [False, ...., False, True, ..., True] and we need to find the first place of True.

Complexity

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

Code

class Solution:
    def kthSmallestSubarraySum(self, nums, k):
        n = len(nums)
        acc = [0] + list(accumulate(nums))
        
        def check(T):   #gives number of sums <= T
            beg, end, ans = 0, 0, 0
            while end < n + 1:
                if acc[end] - acc[beg] <= T:
                    end += 1
                else:
                    ans += n + 1 - end
                    beg += 1
            return n*(n+1)//2 - ans >= k
        
        beg, end = 0, sum(nums)
        while beg + 1 < end:
            mid = (beg + end)//2
            if check(mid):
                end = mid
            else:
                beg = mid
                
        return end