LeetCode 面试题40.最小的k个数

题目

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

示例

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

题目链接

题目分析

我们可以遍历原数组，维护一个大小为k的大顶堆：

• 当大顶堆元素个数小于k时，将数组元素插入堆中；
• 当大顶堆元素等于k时：
  • 如果堆顶元素小于数组当前元素，不做操作；
  • 如果堆顶元素大于数组当前元素，则交换这两个值，重新调整大顶堆;
• 最后返回这个堆即可。

遍历调整的过程如下：

for (int i = 0; i < arrSize; i++){
    if (count < k) {
        res[count++] = arr[i];
        if (count == k) buildHeap(res, count);
    }
    else if (count == k){
        if (res[0] > arr[i]) {
            res[0] = arr[i];
            buildHeap(res, count);
        }
    } 
}

---

题目解答

void swap(int* arr, int i, int j){
    int temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
}

void heapify(int* arr, int size, int i){
    if (i >= size) return ;
    int c1 = 2 * i + 1;
    int c2 = 2 * i + 2;
    int max = i;
    if (c1 < size && arr[c1] > arr[max]) max = c1;
    if (c2 < size && arr[c2] > arr[max]) max = c2;
    if (max != i) {
        swap(arr, max, i);
        heapify(arr, size, max);
    }
}

void buildHeap(int* arr, int size){
    int last = size - 1;
    int parent = (last - 1) / 2;
    for (int i = parent; i >= 0; i--){
        heapify(arr, size, i);
    }
}

int* getLeastNumbers(int* arr, int arrSize, int k, int* returnSize){
    *returnSize = k;
    if (k == 0) return arr;

    int* res = (int*)malloc(sizeof(int) * k);
    int count = 0;
    for (int i = 0; i < arrSize; i++){
        if (count < k) {
            res[count++] = arr[i];
            if (count == k) buildHeap(res, count);
        }
        else if (count == k){
            if (res[0] > arr[i]) {
                res[0] = arr[i];
                buildHeap(res, count);
            }
        } 
    }

    return res;
}