最小的K个数

你快乐吗?

Posted by bbkgl on September 26, 2019

最小的K个数

C++,topK问题。

一般来说这种题有两种种“可行”的方法:

  • 最小/大堆,O(nlogk);
  • 快排的partition思想,O(n);

在C++中,可以用STL中的优先队列来建立最小堆,但面试手撕代码的时候可不会让用库函数和容器。所以这里用最大堆来实解决topK的问题,注意三个测例的坑点:

  • k <= 0,报超时或者段错误;
  • k > input.size(),报超时或者段错误;
  • input为空,段错误。

代码如下

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        if (input.empty() || k == 0 || k > input.size())
            return vector<int> {};
        vector<int> heap(input.begin(), input.begin() + k);
        for (int i = k / 2 - 1; i >= 0; i--)
            adjust(heap, i);
        for (int i = k; i < input.size(); i++) {
            if (input[i] < heap[0]) {
                swap(input[i], heap[0]);
                adjust(heap, 0);
            }
        }
        return heap;
    }
    
    void adjust(vector<int> &heap, int pa) {
        while (pa < heap.size()) {
            int child = pa * 2 + 1;
            if (child >= heap.size()) 
                break;
            if (child + 1 < heap.size() && heap[child+1] > heap[child])
                child++;
            if (heap[pa] < heap[child]) 
                swap(heap[child], heap[pa]);
            else
                break;
            pa = child;
        }
    }
};