题解 | #最小的K个数#

快排思想

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        srand(time(NULL));
        if (input.size() == 1 && k == 1) return input;
        vector<int> ans;
        partition(input, 0, input.size() - 1, k, ans);
        return ans;
    }

    void partition(vector<int>& nums, int start, int end, int k, vector<int>& ans) {
        if (k == 0) return;
        int index = rand() % (end - start + 1) + start;
        swap(nums[start], nums[index]);
        int pivot = nums[start];
        int left = start, right = end;
        while (left < right) {
            while (left < right && nums[right] > pivot) right--;
            if (left < right) nums[left] = nums[right];
            while (left < right && nums[left] <= pivot) left++;
            if (left < right) nums[right] = nums[left];
        }
        nums[left] = pivot;
        if (left - start + 1 <= k) {
            for (int i = start; i <= left; i++) {
                cout << nums[i];
                ans.push_back(nums[i]);
            }
            partition(nums, left+1, end, k - (left - start + 1), ans);
        } else {
            partition(nums, start, left - 1, k, ans);
        }
    }
};