库存管理 III

标签: 数组 分治 快速选择 排序 堆(优先队列)

难度: Easy

仓库管理员以数组 stock 形式记录商品库存表,其中 stock[i] 表示对应商品库存余量。请返回库存余量最少的 cnt 个商品余量,返回 顺序不限

示例 1:

输入:stock = [2,5,7,4], cnt = 1
输出:[2]

示例 2:

输入:stock = [0,2,3,6], cnt = 2
输出:[0,2] 或 [2,0]

提示:

  • 0 <= cnt <= stock.length <= 10000
    0 <= stock[i] <= 10000

Submission

运行时间: 116 ms

内存: 15.8 MB

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return []
        if len(arr) <= k:
            return arr

        self.partition_helper(arr, 0, len(arr) - 1, k)
        return arr[:k]

    def partition_helper(self, arr, lo, hi, k):
        j = self.partition(arr, lo, hi)
        if j > k:
            self.partition_helper(arr, lo, j-1, k)
        elif j < k:
            self.partition_helper(arr, j+1, hi, k)

    def partition(self, arr, lo, hi):
        if lo == hi:
            return lo
        v = arr[lo]
        i = lo
        j = hi+1
        while True:
            while True:
                i += 1
                if arr[i] > v:
                    break
                if i >= hi:
                    break
            while True:
                j -= 1
                if arr[j] < v:
                    break
                if j <= lo:
                    break
            if i >= j:
                break
            arr[i], arr[j] = arr[j], arr[i]
        arr[lo], arr[j] = arr[j], arr[lo]
        return j

Explain

该题解采用了快速选择(Quickselect)算法,这是一种用于在未排序的数组中找到第 k 个最小元素的方法。算法的核心思想是快速排序的一部分,即选取一个枢轴元素并对数组进行分区,将小于枢轴的元素放在其左侧,大于枢轴的元素放在其右侧。根据枢轴的位置与 k 的比较,递归地在需要的一侧继续进行分区操作,直到找到第 k 个最小元素。在本题中,我们不需要完全排序,只需要找到最小的 k 个元素即可。

时间复杂度: O(n)

空间复杂度: O(log n)

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return []  # 如果 k 为 0,直接返回空列表
        if len(arr) <= k:
            return arr  # 如果数组长度不大于 k,直接返回整个数组

        self.partition_helper(arr, 0, len(arr) - 1, k)  # 调用辅助函数来进行分区
        return arr[:k]  # 返回最小的 k 个数

    def partition_helper(self, arr, lo, hi, k):
        j = self.partition(arr, lo, hi)  # 执行一次分区操作
        if j > k:
            self.partition_helper(arr, lo, j-1, k)  # 如果枢轴位置大于 k,递归左侧
        elif j < k:
            self.partition_helper(arr, j+1, hi, k)  # 如果枢轴位置小于 k,递归右侧

    def partition(self, arr, lo, hi):
        if lo == hi:
            return lo
        v = arr[lo]  # 选择枢轴元素
        i = lo
        j = hi+1
        while True:
            while True:
                i += 1
                if i >= hi or arr[i] > v:
                    break
            while True:
                j -= 1
                if j <= lo or arr[j] < v:
                    break
            if i >= j:
                break  # 当两个指针相遇时跳出循环
            arr[i], arr[j] = arr[j], arr[i]  # 交换元素
        arr[lo], arr[j] = arr[j], arr[lo]  # 将枢轴元素放到正确位置
        return j  # 返回枢轴的最终位置

Explore

题解确实处理了 k 等于 0 的情况,返回空列表,这是一个有效的边界条件处理。然而,题解也应该明确处理 k 为负数的情况,因为负数不是有效的 k 值。另外,如果 k 大于数组长度,虽然题解返回整个数组,但这种情况下可能也应该有特定的处理或者提示,以确保函数对所有输入值都进行了适当的错误处理和边界检查。

快速选择算法的递归实现确实存在堆栈溢出的风险,尤其是在最坏情况下,即每次选择的枢轴都不将数组有效分割成较小的部分时。这种情况下递归的深度可以达到数组的长度,从而导致堆栈溢出。为了减少这种风险,可以采用迭代方法,或者改进枢轴选择策略,如使用随机枢轴或三数中值分割法来尽可能确保分割的均匀性。

快速选择算法通过递归地调整枢轴的位置来确保枢轴左侧的所有元素都不大于右侧的元素。当枢轴的最终位置正好是 k-1 时,它及其左侧的所有元素(总共 k 个元素)就是数组中最小的 k 个元素。算法将只在需要的部分(即包含 k 个最小元素的部分)上进行递归操作,从而在递归结束时保证得到正确的最小 k 个元素。

在分区函数中,内部的两个 while 循环通过条件检查来确保 i 和 j 不会越界。对于 i 的循环,它会停止递增当 i 为 hi 或者 arr[i] 大于枢轴值 v。对于 j 的循环,它会停止递减当 j 为 lo 或者 arr[j] 小于 v。这样的条件确保了即使在极端情况下,如数组已经有序或枢轴选取不佳时,循环也不会导致数组访问越界。