题解 | 最小的K个数

class Solution {
public:

    /**
        partition function in quick sort, range [lh, rh)
        returns pivot index after partition
    */
    int Partition(vector<int> &input, int lh, int rh) {
        int pivot = lh;
        int index = pivot + 1;
        auto swap = [](vector<int> &input, int i1, int i2) -> void {
            int tmp = input[i1];
            input[i1] = input[i2];
            input[i2] = tmp;
        };

        for(int i = index; i < rh; ++i) {
            if(input[i] < input[pivot]) {
                swap(input, i, index++);
            }
        }
        swap(input, index - 1, pivot);
        return index - 1;
    }
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * 
     * @param input int整型vector 
     * @param k int整型 
     * @return int整型vector
     */
    vector<int> GetLeastNumbers_Solution(vector<int>& input, int k) {
        if(k == 0) {
            return {};
        }
        int lh = 0, rh = input.size(), p;
        while(lh < rh) {
            p = Partition(input, lh, rh);
            if(p + 1 == k) {
                return vector<int>(input.begin(), input.begin() + k);
            } else if(p < k) {
                lh = p + 1;
            } else {
                rh = p;
            }
        }
        // NEVER REACHED HERE
        return {-1};
    }
};