第k大-二分查找-1139. 第k大的子數組

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;
    }