题解 | 最小的K个数

class Solution {
public:

    /**
        在input中[lh, rh)查找前k小的数并返回
    */
    vector<int> GetLeastNumbers(vector<int> &input, int lh, int rh, int k) {
        if(k == 0) {
            return {};
        }
        if(rh - lh == 1) {
            return {input[lh]};
        }

        int pivot = lh;
        int index = pivot + 1;
        for(int i = index; i < rh; ++i) {
            if(input[i] > input[pivot]) {
                int tmp = input[i];
                input[i] = input[index];
                input[index] = tmp;
                ++ index;
            }
        }

        vector<int> result;
        int rank = rh - index + 1;
        if(rank > k) {
            return GetLeastNumbers(input, index, rh, k);
        } else {
            // rank <= k
            result.push_back(input[pivot]);
            if(rank > 1) {
                auto smaller = GetLeastNumbers(input, index, rh, rank - 1);
                result.insert(result.end(), smaller.begin(), smaller.end());
            }
            if(rank < k) {
                auto larger = GetLeastNumbers(input, pivot + 1, index, k - rank);
                result.insert(result.end(), larger.begin(), larger.end());
            }
        }

        return result;
    }
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * 
     * @param input int整型vector 
     * @param k int整型 
     * @return int整型vector
     */
    vector<int> GetLeastNumbers_Solution(vector<int>& input, int k) {
        // write code here
        return GetLeastNumbers(input, 0, input.size(), k);
    }
};