002-4算法笔记【分治】最小K个数

面试题 17.14. 最小K个数

快速搜索排序：找到第k+1小的数，返回左侧部分
复杂度: O(n) = n +n/2 +n/4 +… < 2*n

进行一轮快排后，若
1.i==k，arr[i]是第k+1小的数，返回左侧部分
2.i<k，第k+1小的数在右侧，搜索[i+1,right]
3.i>k，第k+1小的数在左侧，搜索[left,i-1]

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        if (k >= arr.size()) return arr;
        return quickSort(arr, k, 0, arr.size() - 1);
    }
private:
    //快速搜索排序：找到第k+1小的数，返回左侧部分
    //复杂度:O(n) = n +n/2 +n/4 +... < 2*n
    //进行一轮快排后，若
    //1.i==k，arr[i]是第k+1小的数，返回左侧部分
    //2.i<k，第k+1小的数在右侧，搜索[i+1,right]
    //3.i>k，第k+1小的数在左侧，搜索[left,i-1]
    vector<int> quickSort(vector<int>& arr, int k, int left, int right) {
        int i = left, j = right;
        while (i < j) {
            while (i < j && arr[j] >= arr[left]) j--;
            while (i < j && arr[i] <= arr[left]) i++;
            swap(arr[i], arr[j]);
        }
        swap(arr[i], arr[left]);

        if (i > k) return quickSort(arr, k, left, i - 1);
        if (i < k) return quickSort(arr, k, i + 1, right);
        //i==k, 返回左部
        vector<int> res;
        res.assign(arr.begin(), arr.begin() + k);
        return res;
    }
};