題目
題目鏈接
輸入整數數組 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]
題解
根據數組的前k個數建立一個大頂堆,對,是大頂堆,然後數組後面一次更新這個頂堆。
class Solution {
public:
void sink(int i, int k, vector<int> &res)
{
while ((2 * i + 1) <= (k - 1))
{
int child = 2 * i + 1;
if ((child + 1) < k && res[child] < res[child + 1])
{
child++;
}
if (res[i] >= res[child])
break;
int tmp = res[i];
res[i] = res[child];
res[child] = tmp;
i = child;
}
}
vector<int> getLeastNumbers(vector<int>& arr, int k) {
vector<int> res(k);
if(!k)
return res;
for (int i = 0; i < k; i++)
{
res[i] = arr[i];
}
// 構建節點個數爲k的大頂堆
for (int i = k / 2 - 1; i >= 0; i--)
{
sink(i, k, res);
}
for (unsigned int i = k; i < arr.size(); i++)
{
if (arr[i] < res[0])
{
res[0] = arr[i];
sink(0, k, res);
}
}
return res;
}
};