《算法竞赛进阶指南》0x17二叉堆 POJ2442 矩阵取数求前N大

Posted 2020-12-10 randy-lo

题目链接：http://poj.org/problem?id=2442

给定一个M*N的矩阵，要求从每一行中都取出一个数然后累加，问最小的N个累积和为多少。使用堆可以在O(MNlogN)时间复杂度内求出。

M行的最大取法一定是通过前M-1行的最大取法+第M行取数然后求前N大获取的，所以有归纳法可以考虑两组数据进行选择。

代码：

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<vector>
#include<queue>
using namespace std;
#define maxn 2004
typedef pair<int,int> P;

int a[maxn],b[maxn],c[maxn];
int n,m;
void merge(){//将ab中最小的N对数的和存入a中 
    priority_queue<P,vector<P>,greater<P> > heap;//保存值以及a数组中的下标 
    for(int i=0;i<n;i++)heap.push(make_pair(b[i]+a[0],0));
    for(int i=0;i<n;i++){
        P t=heap.top();
        heap.pop();
        int s=t.first,p=t.second;
        c[i]=s;
        heap.push(make_pair(s-a[p]+a[p+1],p+1));
    }
    for(int i=0;i<n;i++)a[i]=c[i];
}
int main(){
    int T;
    cin>>T;
    while(T--){
        scanf("%d%d",&m,&n);
        for(int i=0;i<n;i++)scanf("%d",&a[i]);
        sort(a,a+n);
        for(int i=0;i<m-1;i++){
            for(int j=0;j<n;j++)scanf("%d",&b[j]);
            merge();
        }
        for(int i=0;i<n;i++)printf("%d ",a[i]);
        puts("");
    }
}