第 K 小的子数组和·

Submission

运行时间: 565 ms

内存: 18.2 MB

class Solution:
    def kthSmallestSubarraySum(self, nums: List[int], k: int) -> int:
        n = len(nums)
        l, r = min(nums), sum(nums)
        while l < r:
            mid = l + r >> 1
            def check(thred):
                cnt = 0
                s = j = 0
                for i in range(n):
                    s += nums[i]
                    while s > thred:
                        s -= nums[j]
                        j += 1
                    cnt += i - j + 1
                return cnt < k
            if check(mid):
                l = mid + 1
            else:
                r = mid
        return l

Explain

此题解采用二分查找结合滑动窗口的方法来找到第 K 小的子数组和。首先,设置搜索区间的左端点 l 为数组中的最小值,右端点 r 为数组元素之和。\n\n在二分查找过程中,计算中点 mid,然后使用辅助函数 check 来判断小于等于 mid 的子数组和的数量是否小于 k。如果是,则说明第 K 小的子数组和大于 mid,需要将搜索区间的左端点调整为 mid+1;否则,将右端点调整为 mid。\n\n该过程持续直到 l 和 r 相遇,此时的 l(或 r)即为第 K 小的子数组和。

时间复杂度: O(n log(sum(nums) - min(nums)))

空间复杂度: O(1)

class Solution:
    def kthSmallestSubarraySum(self, nums: List[int], k: int) -> int:
        n = len(nums)  # 数组长度
        l, r = min(nums), sum(nums)  # 初始化左右边界
        while l < r:
            mid = (l + r) // 2  # 计算中点
            def check(thred):  # 辅助函数,用于判断小于等于 thred 的子数组和的个数是否小于 k
                cnt = 0  # 子数组和的计数
                s = j = 0  # 滑动窗口的总和和起始指针
                for i in range(n):
                    s += nums[i]  # 扩展窗口
                    while s > thred:  # 调整窗口直到窗口内的和不大于 thred
                        s -= nums[j]
                        j += 1
                    cnt += i - j + 1  # 统计符合条件的子数组和的数量
                return cnt < k
            if check(mid):
                l = mid + 1  # 移动左边界
            else:
                r = mid  # 移动右边界
        return l  # 返回第 K 小的子数组和

Explore

选择二分查找方法是为了提高算法的效率。直接计算所有可能的子数组和并进行排序的方法涉及到生成所有子数组和,其数量为O(n^2),并且排序这些和需要O(n^2 log n^2)的时间复杂度,这在数据量较大时会非常低效。相比之下,二分查找方法通过逐步缩小可能的值的范围来查找目标子数组和,结合滑动窗口的方式来统计符合条件的子数组数量,总体时间复杂度可以降低到O(n log S),其中S是数组元素之和的范围,这种方法更加高效。

左边界`l`被设定为数组中的最小值,因为这是单个元素形成的子数组和中的最小可能值,而右边界`r`设定为数组元素之和,因为这是单个子数组(即包含所有元素的子数组)可以达到的最大和。这种设定在非负数数组中总是适用。如果数组包含负数,这种方法仍然适用,因为我们通常不考虑负数数组中和为负值的情况(对于查找第k小的子数组和)。设定范围从最小值到总和能确保包含所有可能的子数组和。

辅助函数`check`采用滑动窗口方法来有效地统计所有小于等于`mid`的子数组和的数量。通过维护一个窗口的总和`s`,当窗口的总和超过`mid`时,通过移动窗口的起始位置`j`来减小总和,直到窗口内的总和不大于`mid`。每次扩展窗口(增加`i`),都可以计算以`i`结尾,以j开始到i之间所有可能的子数组和,这些都是小于等于`mid`的。这种方法避免了枚举所有子数组,从而在每次滑动时以线性时间完成统计,是一种时间效率较高的实现方式。

在二分查找中,如果`check`函数返回`true`,意味着当前的`mid`值作为子数组和的上限,可以找到的符合条件的子数组数量少于k个。这表明第k小的子数组和必须大于`mid`,因此需要将搜索区间的左边界调整为`mid + 1`来排除当前的`mid`值。这种调整是安全的,因为已知`mid`值太小,不能满足第k小的条件。通过这种方法,我们确保不会错过正确答案,而是逐步逼近正确的最小子数组和。