最小的K个数

题目描述

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

 

思路1

维持一个大根堆，使堆中的元素始终保持为K。

C++STL中，priority默认是大根堆

小根堆定义：priority_queue< int ,vector<int>,greater<int >> q

 

思路1代码

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> res;
        int n = input.size();
        if (k > n || k < 1 || n == 0) return res;
        priority_queue<int,vector<int>,less<int>> q;
        for (int i = 0; i < n; ++i){
            q.push(input[i]);
            if (q.size() > k) q.pop();
        }
        for (int i = 0; i < k; ++i){
            res.emplace_back(q.top());
            q.pop();
        }
        return res;
    }
};

 

思路2

利用快排思想。时间复杂度O(n）。

 

思路2代码

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> res;
        int n = input.size();
        if (k < 1 || k > n || n == 0) return res;
        int left = 0, right = n-1;
        int pivot = partition(input, left, right);
        while (pivot != k-1){
            if (pivot > k-1) right = pivot-1;
            else left = pivot + 1;
            pivot = partition(input, left, right);
        }
        for (int i = 0; i < k; ++i)
            res.emplace_back(input[i]);
        return res;
    }
    
    int partition(vector<int>& arr, int left, int right){
        int target = arr[right];
        int pos = left-1;
        for (int i = left; i <= right; ++i){
            if (arr[i] <= target){
                pos++;
                swap(arr[pos], arr[i]);
            }
        }
        return pos;
    }
};