【分治】最小的k个数

剑指 Offer 40. 最小的k个数

剑指 Offer 40. 最小的k个数

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：

• 0 <= k <= arr.length <= 10000
• 0 <= arr[i] <= 10000

---

解题思路：快速选择算法

​ 这道题和 215. 数组中的第K个最大元素 一样，同样可以使用排序或者堆来解决，但是它们的时间复杂度分别是 O(n * logn) 和 O(n * logk)，还达不到这道题的极致，我们可以采用前面学的快速排序中分三块区间的思想来优化这道题，使得时间复杂度接近 O(n)。

​ 还是一样，随机选取数组元素为基准值 + 分三区间 的思想不变，变的只不过是分区间后的处理，所以这里的随机数选取以及分区思想就不再赘述，具体可以参考前面的笔记！比如下图，假设此时是第一次进行分区和分类完成之后的效果：

​ 和 215. 数组中的第K个最大元素 不一样的是，这次要求的是最小的 k 个数，也就是说和上次相反，并且返回的是一个数组，但是思路还是一样的！

​ 如上图所见，各指针的位置其实就确定了各区间元素的个数，我们可以得到 a = left - LEFT + 1，b = right - left - 1，这个每道题都不太一样，要具体问题具体分析！

​ 因为这道题要返回的是数组类型，所以我们可以考虑 在原数组上，将前 k 个最小的数放到最左边，然后直接返回即可，所以此时分区完之后 k 也同样是有三种情况，如下所示：（这道题没要求返回数组要有序）

1. 如果 k < a 的话：

   • 说明前 k 个元素都在 [LEFT, left] 区间中，但是因为 [LEFT, left] 区间中的元素是无序的，所以我们 需要继续递归 [LEFT, left] 区间，寻找前 k 个最小的数，但是后面的两个区间就不需要考虑了！

     • 要注意的是，这里的 k == a 的情况 就相当于前 k 个元素就是 [LEFT, left] 区间的所有元素，所以不需要继续递归，直接交给下一个 if 语句中的 return 来处理即可，所以不需要包括进来！

     

2. 如果 k ≤ a + b 的话：

   • 这种情况就是前 k 个最小的数已经覆盖了第一个全都小于 key 的区间，但是不超过全都等于 key 的区间，又因为 [left + 1, right - 1] 这个区间本质上所有元素都是相等的（也就是一种有序），所以只要 k 的个数到达了该区间但是不超过该区间，则可以得到前 k 个最小的数就是 [LEFT, k] 了，所以 直接返回即可，不需要再递归寻找了！

     

3. 除了上面之外的情况：

   • 这种情况也很简单，因为前面的区间 [LEFT, right - 1] 说明已经是包含在了前 k 个最小的数中了，所以我们只需要 递归去找出剩下的在 [right, RIGHT] 区间中的 k - a - b 个最小的数即可！

     

​ 剩下的问题就交给递归去解决，这就是分治的好处！

class Solution {
public:
    // 选随机数作为基准值
    int partition(vector<int>& arr, int LEFT, int RIGHT) { return arr[(rand() % (RIGHT - LEFT + 1)) + LEFT]; }

    // 快速选择算法
    void quick_select(vector<int>& arr, int LEFT, int RIGHT, int k)
    {
        if(LEFT >= RIGHT)
            return;
        
        int key = partition(arr, LEFT, RIGHT);
        int i = LEFT, left = LEFT - 1, right = RIGHT + 1;
        while(i < right)
        {
            if(arr[i] < key) swap(arr[++left], arr[i++]);
            else if(arr[i] == key) i++;
            else swap(arr[--right], arr[i]);
        }

        // 分别获取左边和中间两个区间的元素个数a和b
        // 然后根据k来分情况讨论
        int a = left - LEFT + 1, b = right - left - 1;

        if(k < a) // 注意如果等于a了，那么直接走下一个if语句中的return语句即可，无需继续递归
            quick_select(arr, LEFT, left, k);
        else if(k <= a + b)
            return;
        else
            quick_select(arr, right, RIGHT, k - a - b);
    }

    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        srand(time(nullptr));
        quick_select(arr, 0, arr.size() - 1, k);
        return { arr.begin(), arr.begin() + k }; // 返回的是原数组的前k区间
    }
};