两个有序数组间相加和的Topk问题

两个有序数组间相加和的Topk问题

题目描述

给定两个有序数组arr1和arr2，再给定一个整数k，返回来自arr1和arr2的两个数相加和最大的前k个，两个数必须分别来自两个数组

按照降序输出

[要求]

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

输入描述：

第一行三个整数N, K分别表示数组arr1, arr2的大小，以及需要询问的数
接下来一行N个整数，表示arr1内的元素
再接下来一行N个整数，表示arr2内的元素

输出描述：

输出K个整数表示答案

示例1

输入

5 4
1 2 3 4 5
3 5 7 9 11

输出

16 15 14 14

备注：

$1⩽N⩽10^5$
$0⩽ar{r}_{1_i},ar{r}_{2_i}⩽10^9$
保证$1⩽K⩽2N$

---

题解：

一般涉及到 TopK 问题，首先应该想到使用堆2333。

我们可以知道 a[n-1] + b[n-1] 是整体的最大值，那么在这个最大值之后应该是哪个元素呢，应该是： m a x { a [ n − 1 ] + b [ n − 2 ] , a [ n − 2 ] + b [ n − 1 ] } max\{a[n-1] + b[n-2], a[n-2] + b[n-1]\} max{a[n−1]+b[n−2],a[n−2]+b[n−1]}，它俩就是下一个可能最大值，绝不会出现比它们更大的元素。

我们可以使用一个大根堆，每个节点保存三个元素 { s u m , i d x 1 , i d x 2 } \{sum, idx1, idx2\} {sum,idx1,idx2}，且 $a[idx1]+b[idx2]=sum$，两个下标用来表明该元素由两个数组中哪两个元素凑成的。

选择堆顶元素，其就是整体的最大值，将其删除，然后把两个可能的最大元素插入到堆中去：

• { a [ i d x 1 − 1 ] + b [ i d x 2 ] , i d x 1 − 1 , i d x 2 } \{a[idx1-1]+b[idx2], idx1-1, idx2\} {a[idx1−1]+b[idx2],idx1−1,idx2}

• { a [ i d x 1 ] + b [ i d x 2 − 1 ] , i d x 1 , i d x 2 − 1 } \{a[idx1] + b[idx2-1], idx1, idx2-1\} {a[idx1]+b[idx2−1],idx1,idx2−1}

这样进行 k 次操作后，就可以得到整体的 TopK 了。

注意：
1. 需要判重，因为两个节点的下标有可能相等
2. 测试数据出现无序的情况。。。。。。

代码：

#include <cstdio>
#include <queue>
#include <algorithm>
#include <unordered_map>

using namespace std;

const int N = 100000;

typedef pair<int, int> PII;
typedef pair<int, PII> PIP;

int n, k;
int a[N], b[N];

int main(void) {
    scanf("%d%d", &n, &k);
    for ( int i = 0; i < n; ++i )
        scanf("%d", a + i);
    for ( int i = 0; i < n; ++i ) 
        scanf("%d", b + i);
    sort( a, a + n );
    sort( b, b + n );
    priority_queue<PIP> heap;
    unordered_map<int, bool> hash;
    heap.push({a[n-1]+b[n-1], {n-1,n-1}});
    hash[(n - 1) * N + n - 1] = true;
    while ( k-- ) {
        auto t = heap.top();
        heap.pop();
        printf("%d ", t.first);
        int x = t.second.first, y = t.second.second;
        int key = x * N + y - 1;
        if ( y && !hash.count(key) ) {
            heap.push({a[x] + b[y-1], {x, y-1}});
            hash[key] = true;
        }
        key = (x - 1) * N + y;
        if ( x  && !hash.count(key) ) {
            heap.push({a[x-1] + b[y], {x-1, y}});
            hash[key] = true;
        }
    }
    return 0;
}

贴一个正儿八经的 O(klogk) 复杂度代码：

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

using namespace std;

const int N = 1e5 + 10;

int a[N];
int b[N];

using PII = pair<int, int>;

int main(void) {
    int n, k;
    scanf("%d%d", &n, &k);
    for ( int i = 0; i < n; ++i )
        scanf("%d", a + i);
    for ( int i = 0; i < n; ++i )
        scanf("%d", b + i);
    sort(a, a + n);
    sort(b, b + n);
    
    auto cmp = [&](const PII& lhs, const PII& rhs) {
        return a[lhs.first] + b[lhs.second] <= a[rhs.first] + b[rhs.second];
    };
    
    priority_queue<PII, vector<PII>, decltype(cmp)> heap(cmp);
    
    int r = max(0, n - k);
    for ( int i = n - 1; i >= r ; --i ) {
        heap.emplace( move(PII(i, n - 1)) );
    }
    while ( k-- ) {
        auto [x, y] = heap.top();
        heap.pop();
        printf("%d ", a[x] + b[y]);
        if ( y ) heap.emplace( move(PII(x, y-1)) );
    }
    return 0;
}