题目描述

1918 第 K 小的子数组和

给你一个 长度为 n 的整型数组 nums 和一个数值 k ，返回 第k小的子数组和。

子数组 是指数组中一个 非空 且不间断的子序列。 子数组和 则指子数组中所有元素的和。

示例 1:

输入: nums = [2,1,3], k = 4
输出: 3
解释: [2,1,3] 的子数组为：
- [2] 和为 2
- [1] 和为 1
- [3] 和为 3
- [2,1] 和为 3
- [1,3] 和为 4
- [2,1,3] 和为 6 
最小子数组和的升序排序为 1, 2, 3, 3, 4, 6。 第 4 小的子数组和为 3 。

示例 2：

输入：nums = [3,3,5,5], k = 7
输出：10
解释：[3,3,5,5] 的子数组为：
- [3] 和为 3
- [3] 和为 3
- [5] 和为 5
- [5] 和为 5
- [3,3] 和为 6
- [3,5] 和为 8
- [5,5] 和为 10
- [3,3,5], 和为 11
- [3,5,5] 和为 13
- [3,3,5,5] 和为 16
最小子数组和的升序排序为 3, 3, 5, 5, 6, 8, 10, 11, 13, 16。第 7 小的子数组和为 10 。

提示:

• n == nums.length
• 1 <= n <= 2 * 10^4
• 1 <= nums[i] <= 5 * 10^4
• 1 <= k <= n * (n + 1) / 2

二分查找

对于长度为 n 的数组 \(\textit{nums}\)，共有 \(\dfrac{n(n+1)}{2}\)个子数组，每个子数组都有对应的子数组和，即共有 \(\dfrac{n(n+1)}{2}\)个子数组和。由于数组 \(\textit{nums}\) 的元素都大于 \(0\)，因此最小的子数组和为数组 \(\textit{nums}\) 的最小元素，最大的子数组和为数组 \(\textit{nums}\) 的所有元素之和。第\(k\)小的子数组和一定大于或等于最小的子数组和且小于或等于最大的子数组和。当阈值越大时，子数组和小于或等于阈值的子数组的数量也越大，因此可以使用二分查找的方法寻找第 \(k\) 小的子数组和。

二分查找的做法是，初始化上界和下界分别为最大的子数组和和最小的子数组和，每次计算上界和下界的平均值 \(\textit{mid}\)，计算子数组和小于或等于\(\textit{mid}\)的子数组的数量，将该数量与\(k\)比较之后调整上下界，直到找到第\(k\)小的子数组和。

因为要返回子数组的和，所以这里用到的二段性为子数组的和：

1. \(min(nums) \le\) 任意子数组的和 \(\le sum(nums)\)
2. \(mid = (min(nums) + sum(nums))/2\)，那么二段性可以表示为：以和\(\le mid\) 将所有子数组分成两部分，如果左半部分个数\(< k\)的话，那么第\(k\)小的子数组和一定\(> mid\)。

滑动窗口

对于特定的阈值\(\textit{threshold}\)，计算子数组和小于或等于\(\textit{threshold}\)的子数组的数量可以通过滑动窗口实现。由于数组 \(\textit{nums}\) 的元素都大于 \(0\)，因此对于区间\([\textit{left},\textit{right}]\) ，其中\(0 \le \textit{left} \le \textit{right} \)，如果子数组 \(\textit{nums}[\textit{left}:\textit{right}]\)的元素和小于或等于 \(\textit{threshold}\)，则对于任意 \(\textit{left} \le i \le \textit{right}\)，子数组 \(\textit{nums}[i:\textit{right}]\)的元素和一定小于或等于\(\textit{threshold}\)

因此，对于每个区间的右边界 \(\textit{right}\)（其中 \(0 \le \textit{right} < n\)），需要找到该区间的左边界的临界值 \(\textit{left}\)（即 \(\textit{left}\) 要尽可能小），使得子数组 \(\textit{nums}[\textit{left}:\textit{right}]\) 的元素和小于或等于\(\textit{threshold}\)，则以\(\textit{right}\)为右边界，且元素和小于或等于\(\textit{threshold}\)的子数组有\(\textit{right}-\textit{left}+1\)个。

初始时，\(\textit{left}=\textit{right}=0\)，然后将右边界 \(\textit{right}\)右移直到达到数组末尾，在右边界 \(\textit{right}\) 右移的过程中，左边界的临界值 \(\textit{left}\) 也在右移，因此 \(\textit{left}\)和\(\textit{right}\) 都只会遍历数组一次，可以在 \(O(n)\)的时间内计算元素和小于或等于 \(\textit{threshold}\) 的子数组的数量。

代码

class Solution {
int countSubarrayLess(vector<int>& nums, int mid){
    int left = 0, right = 0;
    int n = nums.size();
    int sum = 0;
    int count = 0;
    while(right < n){
        sum +=nums[right++];
        while(sum>mid){//当left==right时候，和为0，必定小于mid
            sum -=nums[left++];
        }
        count += right - left;
    }
    return count;      
}
public:
int kthSmallestSubarraySum(vector<int>& nums, int k) {
    int min = nums[0], sum =0;
    for(auto& num:nums){
        if(num<min) min = num;
        sum+=num;
    }
    int low = min, high = sum;
    while(low < high){
        int mid = low + (high - low)/2;
        int count = countSubarrayLess(nums, mid);
        if(count<k){
            low = mid + 1;
        }else{
            high = mid;
        }
    }
    return low;
}


};

复杂度分析

• 时间复杂度：\(O(n \log S)\)，其中 \(n\) 为数组 \(\textit{nums}\) 的长度，\(S\) 为数组 \(\textit{nums}\)的元素和。二分查找的次数是 \(O(\log S)\)次，每次二分查找需要 \(O(n)\) 的时间计算元素和小于或等于特定阈值的子数组的数量。
• 空间复杂度：\(O(1)\)。