【LeetCode/LintCode】 题解丨字节跳动试题：第k大的子数组

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

• 1≤n≤10​^5
• 1≤a​i≤10^​9
• 1≤k≤​n(n+1)/2

在线评测地址：点击此处前往

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。

【题解】

算法

二分+two pointer

算法分析

我们可以看到，题目需要求和第k

k大的子区间，而我们的区间总个数共有n(n+1)/2个，当n为10^5​​时这个数高达10^10级别。我们显然不能暴力的枚举每一个区间和然后排序。

算法思路

我们注意到，所有数字的和不超过10^14，这个范围可以让我们想到使用二分最终的答案进行求解。

二分要求解的问题是：对于给定的和x，求有多少个区间的和小于x，小于等于x。这需要我们在O(n)的时间复杂度内进行求解。由于数组内所有数都是正数，我们自然的可以想到同向双指针求解。当当前区间的和大于k，就移动左指针，否则移动右指针。

时间复杂度

O(nlog(n))

public class Solution {
    /**
     * @param a: an array
     * @param k: the kth
     * @return: return the kth subarray
     */

    private int check(long x, int[] a, long k) {
        long tmp1 = 0, tmp2 = 0, now = a[0];
        int l = -1, r = 0, n = a.length;
        long all = (long)n * (n + 1) / 2;
        while (l <= r && r < n)
        {
            if (now >= x) {
                if (now == x) {
                    tmp2++;
                } else {
                    tmp1++;
                }
                tmp1 += n - r - 1;
                l++; now -= a[l];
            }
            else {r++; if (r < n) now += a[r];}
        }
        if (all - tmp1 - tmp2 < k && all - tmp1 >= k) return 0;
        if (all - tmp1 - tmp2 >= k) return 1;
        else return -1;
    }
    public long thekthSubarray(int[] a, long k) {
        // wrrite your code here
        int n = a.length;
        long sum = 0;
        for (int i = 0; i < n; i++) {
            sum += a[i];
        }
        long l = 1, r = sum;
        while (l <= r) {
            long mid = (l + r) / 2;
            int flag = check(mid, a, k);
            if (flag == 0) {
                return mid;
            }
            if (flag == 1) {
                r = mid - 1;
            }
            else {
                l = mid + 1;
            }
        }
        return -1;
    }
}

更多题解参见：

九章算法官网