一、题目描述

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

示例 1:

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

示例 2:

  • 输入:arr = [0,1,2,1], k = 1
  • 输出:[0]

限制:

  • 0 <= k <= arr.length <= 10000
  • 0 <= arr[i] <= 10000

二、题解

解法一:排序

一个最简单的思路是先对所有元素排序,然后输出前面k个元素。使用快速排序的话,时间复杂度O(nlog(n))。

解法二:时间复杂度为O(n)的快速排序

快速排序的特点是:每次排序后,标杆元素左边的元素都比它小,右边的元素都比它大。因此可以利用这一个特性,对数组中的元素进行部分排序,直到返回的标杆元素下标等于k,那么这个标杆元素的左边就是所有需要的目标数据。

相较于快排来说,只需要对目标区间(标杆元素左边或者右边)进行排序就行了,无需对左右两边的区间都进行排序。因此时间复杂度降为O(n)。

解法三:最大堆

最大堆的性质是:堆上层的元素都大于等于下层元素。因此,只要维护一个大小为k的最大堆,循环遍历数组,如果元素大于堆顶元素,就替换掉堆顶元素,重新建堆。堆里面保存的就是k个最小的元素了。时间复杂度O(nlog(k)),空间复杂度O(k)。

三、代码:

3.1 基于快排思想的算法

class Solution {
private:
    int partition(vector<int> &arr, int left, int right) {
        int i, j, key;

        if (left >= right)
            return left;

        i = left;
        j = right;

        key = arr[i];
        while (i < j) {
            while (i < j && arr[j] >= key) {
                j--;
            }
            if (i < j && arr[j] < key) {
                arr[i] = arr[j];
            }

            while (i < j && arr[i] <= key) {
                i++;
            }
            if (i < j && arr[i] > key) {
                arr[j] = arr[i];
            }
        }
        arr[i] = key;

        return i;
    }

public:
    vector<int> getLeastNumbers(vector<int> &arr, int k) {
        int idx, left, right;
        vector<int> output;

        if (arr.size() <= k)
            return arr;
        if (arr.size() == 0)
            return output;

        left = 0;
        right = arr.size() - 1;
        // 循环分区,找到第k个标杆元素
        idx = partition(arr, left, right);
        while (idx != k) {
            if (idx > k) {
                // 标杆元素的下标大于k,要在左区间继续找
                right = idx - 1;
                idx = partition(arr, left, right);
            } else {
                // 标杆元素的下标小于k,要在右区间继续找
                left = idx + 1;
                idx = partition(arr, left, right);
            }
        }

        // 预分配空间,不要直接push_back,否则会浪费很多重新开辟空间的时间
        output.reserve(k);
        for (idx = 0; idx < k; idx++) {
            output.push_back(arr[idx]);
        }

        return output;
    }
};

3.2 使用最大堆

class Solution {
private:
    // 执行最大堆化
    void maxHeapify(vector<int> &v, int idx) {
        int mid, i, left, right, largest;
        left = (idx * 2) + 1;
        right = left + 1;

        if (left < v.size() && v[left] > v[idx])
            largest = left;
        else
            largest = idx;

        if (right < v.size() && v[right] > v[largest])
            largest = right;

        if (largest != idx) {
            swap(v[largest], v[idx]);
            maxHeapify(v, largest);
        }
    }
public:
    vector<int> getLeastNumbers(vector<int> &arr, int k) {
        int i, j, len;
        vector<int> output;

        if (k >= arr.size())
            return arr;

        if (k == 0)
            return output;

        len = 0;
        output.reserve(k);

        for (i = 0; i < arr.size(); i++) {
            if (len < k) {
                // 当最大堆的元素个数小于k时,直接插入到堆
                output.push_back(arr[i]);
                len++;

                // 堆塞满后,构造最大堆
                if (len == k) {
                    for (j = k / 2; j >= 0; j--) {
                        maxHeapify(output, j);
                    }
                }

            } else {
                // 元素小于堆顶元素,替换掉首元素
                if (arr[i] < output[0]) {
                    output[0] = arr[i];
                    maxHeapify(output, 0);
                }
            }
        }

        return output;
    }
};

最后修改:2020 年 03 月 20 日
喜欢就给我点赞吧