剑指offer40 最小的k个数

发表于 | 分类于 |
字数统计: 553 | 阅读时长 ≈ 2

链接

https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/

题意

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

解法

同经典的Top k类似，这里是需要找出最小的k个数
使用分治和堆都可以进行求解
分治会用到快速排序中的Partition方法，对数组进行划分
使得左半边的数都比中间某个数小，右半边的数都比这个数大
直到这个中间的数正好是第k个数 结束运行
时间复杂度为O(n)

另一种解法是使用最小堆
每次保存最小的k个数
如果堆的容量已满 比较堆中最大的数和当前数
较小者放入堆中
时间复杂度为O(nlogk)
它的优点是第一没有修改数组输入的数据。
其二因为内存的大小有限，如果数据量太大不能一次性将数据全部读入
而这一方法只需要维护k个数的堆，适合海量数据存储

代码

分治解法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
public:
    int partition(vector<int>& data, int start, int end) {
        swap(data[start], data[end]);
        int small = start;
        for (int index = start; index < end; index++) {
            if (data[index] < data[end]) {
                swap(data[index], data[small]);
                small++;
            }
        }
        swap(data[small], data[end]);
        return small;
    }
    void solve(vector<int>& arr, int start, int end, int k) {
        if (k == 0) return;
        int count = 0, index = partition(arr, start, end);
        while(index != k - 1) {
            if (index > k - 1) {
                end = index - 1;
                index = partition(arr, start, end);
            }
            else {
                start = index + 1;
                index = partition(arr, start, end);
            }
        }
    }
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        solve(arr, 0, arr.size() - 1, k);
        vector<int> res;
        for (int i = 0; i < k; i++) {
            res.push_back(arr[i]);
        }
        return res;
    }
};

最小堆

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> res;
        if (k == 0) return res;
        priority_queue<int> pq;
        for (int i: arr) {
            if (pq.size() < k) {
                pq.push(i);
            }
            else {
                if (pq.top() > i) {
                    pq.pop();
                    pq.push(i);
                }
            }
        }
        for (int i = 0; i < k; i++) {
            int now = pq.top();
            pq.pop();
            res.push_back(now);
        }
        return res;
    }
};