（算法）前K大的和

题目：

1、有两个数组A和B，每个数组有k个数，从两个数组中各取一个数加起来可以组成k*k个和，求这些和中的前k大。

2、有N个数组，每个数组有k个数，从N个数组中各取一个数加起来可以组成k^N个和，求这些和中的前k大。

思路：

1、将A和B两个数组，按照从大到小排序，那么很容易得到下面的求和矩阵，假设为C，仔细一看，貌似有点规律。

C[0][0]=A[0]+B[0]肯定是最大的，那么候选的第二大的为max(C[0][1],C[1][0])。

我们通过堆来实现，每次从堆中找出最大值C[i][j]，然后把C[i+1][j]和C[i][j+1]加入堆中，直至找到k个最大数。

2、先对前两个数组求前k大和，将结果与第三个数组求前k大和，然后第四个……直到第N个。

代码：

#include<iostream>
#include<vector>
#include<queue>
#include<algorithm>

using namespace std;

struct node{
    int i;
    int j;
    int val;
    node(int x,int y,int v):i(x),j(y),val(v){};
/*
    int operator<(node m) const{
        return val-m.val;
    }
    */
};

struct cmp{
    bool operator()(const node &a,const node &b)const{
        return a.val<b.val;
    }
};

void getKthSum(const vector<int> &arr1,const vector<int> &arr2,int k,vector<int> &result){
    int sz1=arr1.size();
    int sz2=arr2.size();

    if(sz1==0 || sz2==0){
        return;
    }
    
    vector<vector<bool> > visited(sz1,vector<bool>(sz2,false));
    priority_queue<node,vector<node>,cmp> pq;
    int sum=arr1[0]+arr2[0];
    pq.push(node(0,0,sum));
    visited[0][0]=true;
    //result.push_back(sum);
    int count=0;

    while(!pq.empty() && count<k){
        node e=pq.top();
        pq.pop();
        result.push_back(e.val);

        count++;
        int ex1=e.i+1;
        int ey1=e.j;
        int ex2=e.i;
        int ey2=e.j+1;
        if(ex1<sz1 && !visited[ex1][ey1]){
            pq.push(node(ex1,ey1,arr1[ex1]+arr2[ey1]));
            visited[ex1][ey1]=true;
        }
        if(ey2<sz2 && !visited[ex2][ey2]){
            pq.push(node(ex2,ey2,arr1[ex2]+arr2[ey2]));
            visited[ex2][ey2]=true;
        }
    }
}

int main(){
    int m,n,k;
    while(cin>>m>>n>>k){
        vector<int> result; 
        vector<int> arr1(m);
        vector<int> arr2(n);

        for(int i=0;i<m;i++)
            cin>>arr1[i];
        for(int i=0;i<n;i++)
            cin>>arr2[i];
        sort(arr1.begin(),arr1.end(),greater<int>());
        sort(arr2.begin(),arr2.end(),greater<int>());
        getKthSum(arr1,arr2,k,result);

        for(int k=0;k<result.size();k++)
            cout<<result[k]<<" ";
        cout<<endl;
    }
    return 0;
}