n个数组元素最大和的Topk

Posted on | Edited on

两个数组的问题

对于两个递减数组A, B, 分别有m, n个元素，共有m * n个和，我们想求前k大的和，我们可以列出所有的和。

A0B0 > A0B1 > A0B2 > A0B3 > A0B4 ... A1B0 > A1B1 > A1B2 > A1B3 > A1B4 ... A2B0 > A2B1 > A2B2 > A2B3 > A2B4 ... ...

我们把第一行压入优先队列，然后优先队列的第一个一定是和最大的，A0B0，弹出他。然后把A0B0下面的数A1B0压入优先队列，我们可以保证第二大的一定在优先队列中，因为优先队列里面的值都是当列中最大的一个，从而循环即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Sum
{
    public:
        int val;
        int Ai;
        int Bi;
        friend bool operator < (Sum a, Sum b)
        {
            return a.val < b.val;
        }
};

vector<int> maxSum(vector<int> A, vector<int> B, int n)
{
    vector<int> res;
    if( A.size() == 0 )
    {
        return res;
    }
    priority_queue<Sum, vector<Sum>, less<Sum> > qu;
    for( int i = 0 ; i < B.size() ; i++ )
    {
        Sum tmp;
        tmp.val = A[0] + B[i];
        tmp.Ai = 0;
        tmp.Bi = i;
        qu.push(tmp);
    }
    while( res.size() < n )
    {
        res.push_back(qu.top().val);
        Sum tmp = qu.top();
        qu.pop();
        if( ++tmp.Ai < A.size() )
        {
            tmp.val = A[tmp.Ai] + B[tmp.Bi];
            qu.push(tmp);
        }
    }
    return res;
}

对于n个数组

我们先对前两个取前k大的数，在和下面的合并取最大的数，直到取完第n个数组即可。