Leetcode/Lintcode

Lintcode 1139. the kth subarray

Description

Given an array a whose length is n. Return the kth subarray ranked by sums of subarrays.
1 <= n <= 10^5
1 <= a <= 10^9

Solution

The maximum sum of a subarray is bounded by n * a which is 10^14, so we can use a binary search approach to find the kth subarray. Now the problem is how to check the mid of each binary search iteration.

I use an approach similar to prefix sum. The program maintains three variables left, right, and cur_sum(similar to prefix[right]-prefix[left]).

On each iteration, the right index is increased to the right most index which that makes the cur_sum lower or equal to our target, and we can know that there are len(a)-right subarrays that sum greater than our target. The left is then increased by one and the next iteration starts.

Here is my code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Solution:
    """
    @param a: an array
    @param k: the kth
    @return: return the kth subarray
    """
    def rank(self, a, x): # x: target
        n = len(a)
        res = 0
        left = 0
        right = 0
        cur_sum = a[left]
        
        # count the number of subarrays that sum higher than x
        while left < n:
            while cur_sum <= x:
                right += 1
                if (right >= n):
                    break
                cur_sum += a[right]
            if right < n:
                res += n - right
            cur_sum -= a[left]
            left += 1
            
        return (n + 1) * n // 2 - res
        
    def thekthSubarray(self, a, k):
        # wrrite your code here
        low = 0
        high = sum(a)
        while low < high-1:
            mid = low + (high-low)//2
            mid_rank = self.rank(a, mid)
            
            if(mid_rank >= k):
                high = mid
            else:
                low = mid
        return high

Analysis

This solution is O(n * log(n)), where the binary search part runs log(n) times rank function( O(n) ).