打印N个数组整体最大的Top K

打印N个数组整体最大的Top K

题目描述

有N个长度不一的数组，所有的数组都是有序的，请从大到小打印这N个数组整体最大的前K个数。

例如，输入含有N行元素的二维数组可以代表N个一维数组。

219, 405, 538, 845, 971

148, 558

52, 99, 348, 691

再输入整数k=5，则打印：

Top 5: 971, 845, 691, 558, 538

[要求]

时间复杂度为$O(k\log k)$，空间复杂度为$O(k\log k)$

输入描述：

第一行两个整数T, K。分别表示数组个数，需要打印前K大的元素
接下来T行，每行输入格式如下：
开头一个整数N，表示该数组的大小，接下来N个已排好序的数表示数组内的数

输出描述：

从大到小输出输出K个整数，表示前K大。

示例1

输入

3 5
5 219 405 538 845 971
2 148 558
4 52 99 348 691

输出

971 845 691 558 538

备注：

$1⩽\text{数组内数的总个数}⩽10^5$
$0⩽\text{数组内元素}⩽10^9$

保证各个输入值合法

---

题解：

利用堆来解决。具体思路如下：

• 构建一个大小为 N 的大根堆，初始时把每个数组中的最后一个元素加入堆中；
• 建好堆后，堆顶元素就是所有元素中最大的那个，直接打印即可；
• 假设堆顶元素来自第 i 个数组（需要记录堆中每个元素来自哪个数组），那么用第 i 个数组前一个元素替换堆顶元素。若堆顶元素是该数组最后一个元素，则把堆中最后一个元素放在堆顶，堆元素减一。不要忘了更换后调整堆；
• 打印 K 次即可得到 Top K 元素。

PS：优先队列可以直接使用，省很多事。。。但是锻炼编码能力的话，老老实实写堆吧23333。

代码：

#include <cstdio>
#include <vector>

using namespace std;

vector<vector<int> > a;
vector<int> key, val;

int n, m, k, sze;

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

void solve() {
    val.resize( n + 1 );
    key.resize( n + 1 );
    for ( int i = 1; i <= n; ++i ) {
        val[i] = i;
        key[i] = a[i][a[i][0]];
    }
    sze = n;
    for ( int i = n >> 1; i; --i ) down( i );
    while ( k-- ) {
        if ( !sze ) break;
        printf("%d ", key[1]);
        int &t = a[val[1]][0];
        if ( t ) {
            key[1] = a[val[1]][--t];
        } else {
            key[1] = key[--sze];
            val[1] = val[sze];
        }
        down( 1 );
    }
}

int main( void ) {
    scanf("%d%d", &n, &k);
    a.resize( n + 1 );
    for ( int i = 1; i <= n; ++i ) {
        scanf("%d", &m);
        a[i].resize( m + 1);
        a[i][0] = m;
        for ( int j = 1; j <= m; ++j ) {
            scanf("%d", &a[i][j]);
        }
    }
    solve();
    return 0;
}