【剑指offer】面试题40.最小的k个数

2020.7.15更新

分割算法O(n)

class Solution {
public:
    std::vector<int> getLeastNumbers(std::vector<int>& arr, int k) {
        if (k == 0 || arr.size() == 0) return std::vector<int>();
        int s = 0;
        int e = arr.size() - 1;
        int index = Partition(arr, s, e);
        while (index != k - 1) {
            if (index > k - 1) {
                e = index - 1;//注意每次要修改e,s
                index = Partition(arr, s, e);
            }
            else
            {
                s = index + 1;
                index = Partition(arr, s, e);
            }
        }
        return std::vector<int>(arr.begin(), arr.begin() + k);
    }
    int Partition(std::vector<int>& arr, int s, int e) {
        int index = (rand() % (e - s + 1)) + s;//随机获取[s,e]的位置
        std::swap(arr[index], arr[s]);//交换
        int i = s + 1, j = e;
        while (true) {
            while (i <= j && arr[i] <= arr[s]) i++;
            while (i <= j && arr[j] >= arr[s]) j--;
            if (i > j) break;
            std::swap(arr[i], arr[j]);
        }
        std::swap(arr[j], arr[s]);
        return j;
    }
};

快排切分时间复杂度分析： 因为我们是要找下标为k的元素，第一次切分的时候需要遍历整个数组 (0 ~ n) 找到了下标是 j 的元素，假如 k 比 j 小的话，那么我们下次切分只要遍历数组 (0~k-1)的元素就行啦，反之如果 k 比 j 大的话，那下次切分只要遍历数组 (k+1～n) 的元素就行啦，总之可以看作每次调用 partition 遍历的元素数目都是上一次遍历的 1/2，因此时间复杂度是 N + N/2 + N/4 + … + N/N = 2N, 因此时间复杂度是 O(N)O(N)。

大顶堆算法O(nlogk)

class Solution {
public:
    std::vector<int> getLeastNumbers(std::vector<int>& arr, int k) {
        //std::priority_queue<int, std::vector<int>, std::greater<int>> q2;//小顶堆
        std::priority_queue<int> q;
        for (auto e : arr) {
            q.push(e);
            if (q.size() > k) {//如果数量超过k个，弹出堆顶
                q.pop();
            }
        }
        std::vector<int> ans;
        while (!q.empty())
        {
            ans.push_back(q.top());
            q.pop();
        }
        return ans;
    }
};

本题是求前 K 小，因此用一个容量为 K 的大根堆，每次 poll 出最大的数，那堆中保留的就是前 K 小啦（注意不是小根堆！小根堆的话需要把全部的元素都入堆，那是 O(NlogN)😂，就不是 O(NlogK)～～）

解题思路

维护一个k大小的大顶堆，如果当前堆的size>k，则弹出

代码

class Solution
{
public:
	vector<int> getLeastNumbers(vector<int>& arr, int k)
	{
		priority_queue<int, vector<int>, less<int>> q;//大顶堆
        //priority_queue<int, vector<int>, greater<int>> q;//小顶堆
		for (vector<int>::iterator it = arr.begin(); it != arr.end(); it++)
		{
		    q.push(*it);//堆的size<=k，则加入到堆中
			if(q.size()>k) q.pop();
		}
		vector<int> res;
		while (!q.empty())
		{
			res.push_back(q.top());
			q.pop();
		}
		return res;
	}
};