第k大-二分查找-1139. 第k大的子数组

Posted 2021-02-21 hyserendipity

2020-04-25 18:19:22

问题描述：

给定一个长度为n的数组a，它有n * (n + 1) / 2个子数组。请计算这些子数组的和，然后按照升序排列，并返回排序后第k个数。

样例

Example1

Input: 
[2,3,1,4]
6
Output:5
Explanation:
我们可以得到所有子数组的和是 [1,2,3,4,4(3 + 1), 5(1 + 4), 5(2 + 3), 6(2 + 3 + 1), 8(3 + 1 + 4), 10(2 + 3 + 1 + 4)]。其中第六个是5。

注意事项

• 1≤n≤1e?5??
• 1≤a?i??≤1e9??
• 1≤k≤???n(n+1)/2??

问题求解：

Top-k问题最经典的解法是使用优先队列求解，但如果直接使用优先队列，其时间复杂度在本题中是O(n ^ 2)的，肯定会TLE。

另外，如果考虑到以每个起点的序列是有序的，可以直接使用堆来对每个序列进行维护，时间复杂度是O(k)，由于本题中k ~= n ^ 2，所以依然会超时。

正确的解法就是第三种使用二分查找的方式求解。

时间复杂度：O(nlogn)

    public long thekthSubarray(int[] a, long k) {
        long l = 0;
        long r = 0;
        for (int num : a) r += num;
        while (r - l > 1) {
            long mid = l + (r - l) / 2;
            if (helper(a, mid) >= k) r = mid;
            else l = mid;
        }
        return r;
    }
    
    private long helper(int[] nums, long k) {
        int n = nums.length;
        long res = 0;
        int start = 0;
        long sum = 0;
        for (int end = 0; end < n; end++) {
            sum += nums[end];
            if (sum > k) {
                res += n - end;
                sum -= nums[start++];
                while (sum > k) sum -= nums[end--];
            }
        }
        return (long)n * (n + 1) / 2 - res;
    }