Day16:最小的k个数

仓库管理员以数组 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]

LCR 159. 库存管理 III - 力扣（LeetCode）

首先考虑用TreeSet来实现这个代码，因为TreeSet会基于红黑树自动帮我们排序。

class Solution {
    public int[] inventoryManagement(int[] stock, int cnt) {
        TreeSet<Integer> treeSet = new TreeSet<>();
        for(int i = 0; i < stock.length; i++){
            treeSet.add(stock[i]);
        }

        // 取出最小的 cnt 个元素
        int[] result = new int[cnt];
        int index = 0;
        for (int num : treeSet) {
            if (index < cnt) {
                result[index++] = num;
            } else {
                break; // 已经取够 cnt 个元素，退出循环
            }
        }

        return result;
    }
}

很明显，不合适，因为Set集合会去重。

下面改用快排，快排的时间复杂度为O(n)，刚好复习一下快排的代码。

class Solution {
    public int[] inventoryManagement(int[] stock, int cnt) {
        quickSort(stock, 0, stock.length - 1);
        int[] result = new int[cnt];
        for(int i = 0; i < cnt; i++){
            result[i] = stock[i];
        }
        return result;
    }

    public void quickSort(int[] stock, int low, int high){
         if (low < high) {
            // 找到分区点
            int partitionIndex = partition(stock, low, high);

            // 递归排序左半部分
            quickSort(stock, low, partitionIndex - 1);

            // 递归排序右半部分
            quickSort(stock, partitionIndex + 1, high);
        }
    }

    public int partition(int[] stock, int low, int high){
        //找到基准元素
        int pivot = stock[low];
        int left = low + 1;   //左指针
        int right = high;  //右指针
        while(left <= right){
            while(left <= right && stock[right] > pivot){
                right--;
            }
            while(left <= right && stock[left] < pivot){
                left++;
            }

            if(left <= right){
                swap(stock,left,right);
                left++;
                right--;
            }
        }
        swap(stock,right,low);
        return right;
    }

    private void swap(int[] stock, int i, int j) {
        int temp = stock[i];
        stock[i] = stock[j];
        stock[j] = temp;
    }

}

快排的核心全在partition算法里，本质是确定分区点，每一次分区就代表这个元素被排好了。

我们分析一下改怎么写，如何确定最后return的是左边还是右边。

我们把最左端定为哨兵，也就是说最后的位置左边必须比哨兵小，右边必须比哨兵大。while循环的条件是left<=right。首先收缩边界，然后交换，最后的情况是right指着最后一个小于或等于 pivot 的元素。因此要pivot和right换。