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

题目

• 两个有序数组间相加和的Topk问题
• 考查点：堆 + 哈希表

分析

• 数据是十万，所以考虑$O(NlogN)$以下的解法。因为是TopK问题，所以暗示用堆，每次从堆中取出最大值，所以需要大根堆。放入堆中的元素不能有重复，因此要使用哈希表判重。另外，最大值只能是两种情况。
  • a数组往前移动一格，a[i-1] + b[j]
  • b数组往前移动一格，a[i] + b[j-1]
• 所以每次将这两种值扔进大根堆中，取出k个值即为答案。

代码

#include <bits/stdc++.h>
using namespace std;
const int N = 1e5+10;
int a[N], b[N];
using LL = long long;
using VT = vector<LL>;
int main(){
    
    
	int n, k;
	cin>>n>>k;
	for(int i = 0; i < n; i++) cin>>a[i];
	for(int i = 0; i < n; i++) cin>>b[i];
	sort(a, a+n);
    sort(b, b+n);
	priority_queue<VT> pq;
    pq.push({
    
    a[n-1]+b[n-1],n-1,n-1});
    unordered_set<pair<int, int>> st;
    while(k--){
    
    
        auto t =pq.top();
		pq.pop();
        int sum = t[0], i=t[1],j=t[2];
        printf("%d ",sum);
        if(i && !st.count({
    
    i-1,j})){
    
    
            st.insert({
    
    i-1, j});
            pq.push({
    
    a[i-1]+b[j],i-1,j});
        }
        if(j && !st.count({
    
    i,j-1})){
    
    
            st.insert({
    
    i, j-1});
            pq.push({
    
    a[i]+b[j-1],i,j-1});
        }
    }
	return 0;
} 

类似题

• 题目：378. 有序矩阵中第 K 小的元素
• 考查点：堆+哈希表

代码

• 同样的思路

class Solution {
    
    
public:
    using VT = vector<int>;
    int kthSmallest(vector<vector<int>>& mat, int k) {
    
    
        unordered_set<int> st;
        int n = mat.size();
        priority_queue<VT, vector<VT>, greater<VT>> pq;
        pq.push({
    
    mat[0][0], 0, 0});
        st.insert(0);
        while(pq.size()){
    
    
            auto t = pq.top();
            int sum = t[0], i = t[1], j = t[2];
            pq.pop();
            k--;
            // cout<<sum<<endl;
            if(!k) return sum;
            if(i + 1 < n && !st.count((i+1)*n+j)){
    
    
                st.insert((i+1)*n+j);
                pq.push({
    
    mat[i+1][j], i+1, j});
            }
            
            if(j + 1 < n && !st.count(i*n+j+1)){
    
    
                st.insert(i*n+j+1);
                pq.push({
    
    mat[i][j+1], i, j+1});
            }
        }
        return -1;
    }
};