输入整数数组 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个元素。
class Solution {
public:
vector<int> getLeastNumbers(vector<int>& arr, int k) {
vector<int> vec;
sort(arr.begin(), arr.end());
for (int i = 0; i < k; ++i) {
vec.push_back(arr[i]);
}
return vec;
}
};
快排法(参考力扣官方题解)
思路:快排的划分函数每次执行完后都能将数组分成两个部分,小于等于分界值 pivot 的元素的都会被放到数组的左边,大于的都会被放到数组的右边,然后返回分界值的下标。与快速排序不同的是,快速排序会根据分界值的下标递归处理划分的两侧,而这里只需要处理划分的一边。
定义函数 randomized_selected(arr, l, r, k) 表示划分数组 arr 的 [l,r] 部分,使前 k 小的数在数组的左侧,在函数里调用快排的划分函数,假设划分函数返回的下标是 pos(表示分界值 pivot 最终在数组中的位置),即 pivot 是数组中第 pos - l + 1 小的数,那么一共会有三种情况:
如果 pos - l + 1 == k,表示 pivot 就是第 kk 小的数,直接返回即可;
如果 pos - l + 1 < k,表示第 kk 小的数在 pivot 的右侧,因此递归调用 randomized_selected(arr, pos + 1, r, k - (pos - l + 1));
如果 pos - l + 1 > k,表示第 kk 小的数在 pivot 的左侧,递归调用 randomized_selected(arr, l, pos - 1, k)。
函数递归入口为 randomized_selected(arr, 0, arr.length - 1, k)。在函数返回后,将前 k 个数放入答案数组返回即可。
class Solution {
int partition(vector<int>& nums, int l, int r) {
int pivot = nums[r];
int i = l - 1;
for (int j = l; j <= r - 1; ++j) {
if (nums[j] <= pivot) {
i = i + 1;
swap(nums[i], nums[j]);
}
}
swap(nums[i + 1], nums[r]);
return i + 1;
}
// 基于随机的划分
int randomized_partition(vector<int>& nums, int l, int r) {
int i = rand() % (r - l + 1) + l;
swap(nums[r], nums[i]);
return partition(nums, l, r);
}
void randomized_selected(vector<int>& arr, int l, int r, int k) {
if (l >= r) return;
int pos = randomized_partition(arr, l, r);
int num = pos - l + 1;
if (k == num) return;
else if (k < num) randomized_selected(arr, l, pos - 1, k);
else randomized_selected(arr, pos + 1, r, k - num);
}
public:
vector<int> getLeastNumbers(vector<int>& arr, int k) {
srand((unsigned)time(NULL));
randomized_selected(arr, 0, (int)arr.size() - 1, k);
vector<int>vec;
for (int i = 0; i < k; ++i) vec.push_back(arr[i]);
return vec;
}
};