找到无序数组中最小的k个数

找到无序数组中最小的k个数

题目描述

给定一个整型数组arr，找到其中最小的k个数。

输入描述：

输入包含两行，第一行包含两个整数n和k$(1\le k\le n\le 10^5)$，代表数组arr的长度，第二行包含n个整数，代表数组$arr(1\le ar{r}_i\le 10^9)$。

输出描述：

输出包含一行，k个整数，代表数组中最小的k个整数。

示例1

输入

5 3
3 5 1 5 2

输出

3 1 2

备注：

时间复杂度$O(nlo{g}_{2}k)$和$O(n)$，额外空间复杂度$O(n)$。

---

题解：

如果直接对原数组排序，复杂度为$O(nlo{g}_{2}n)$，与题目要求不符。我们可以考虑维护一个大小为 k 的大根堆，对于一个新的元素，若其小于堆顶元素，则将堆顶元素替换为该元素，并向下调整堆。这样遍历完数组中所有元素后，可以得到最小的 k 个数。时间复杂度为$O(nlo{g}_{2}k)$。$O(n)$的解法很麻烦，暂时跳过。。。

代码：

#include <cstdio>
#include <algorithm>

using namespace std;

const int N = 100010;

int n, k;
int h[N];

void down( int u ) {
    int v = u;
    if ( u*2 <= k && h[u*2] > h[u] ) v = u << 1;
    if ( u*2+1 <= k && h[u*2+1] > h[v] ) v = (u << 1) + 1;
    if ( v != u ) {
        swap( h[u], h[v] );
        down( v );
    }
}

int main(void) {
    scanf("%d%d", &n, &k);
    for ( int i = 1; i <= k; ++i ) scanf("%d", h + i);
    for ( int i = k >> 1; i; --i ) down( i );
    int val;
    for ( int i = k + 1; i <= n; ++i ) {
        scanf("%d", &val);
        if ( val < h[1] ) {
            h[1] = val;
            down( 1 );
        }
    }
    for ( int i = 1; i <= k; ++i )
        printf("%d%c", h[i], " \n"[ i == k ]);
    return 0;
}