知识点

堆 多路归并

思路

如果直接维护一个堆来维护最大的k组数对,遍历每一对数,时间复杂度是O(n^2logk),显然超时,不做讨论。

我们考虑把每个priceA的位置当做一个组,然后每个组的起点位于priceB的末尾,建立一个堆选出那个组是最大的。

首先我们把每个priceA的位置和priceB的末尾元素组成的组入堆。之后每次选出当前的最大元素后,我们看这一组是不是后继有人(意思是能不能把该组对应的priceB的指针往前走),如有把该候选数对入堆。因为每次出堆一个,入堆最多一个,堆的大小最多也就和priceA的大小一样。

时间复杂度为O(nlogn + klogn)

AC Code(C++)

#include <queue>
#define x first
#define y second
class Solution {
public:
    /**
     * 代码中的类名、方法名、参数名已经指定,请勿修改,直接返回方法规定的值即可
     *
     * 
     * @param pricesA int整型vector 
     * @param pricesB int整型vector 
     * @param k int整型 
     * @return int整型vector<vector<>>
     */
    using pii = pair<int, int>;
    vector<vector<int> > kPairsWithLargestSums(vector<int>& pricesA, vector<int>& pricesB, int k) {
        int n = pricesA.size(), m = pricesB.size();
        auto cmp = [&](pii& a, pii& b) {
            if (pricesA[a.x] + pricesB[a.y] == pricesA[b.x] + pricesB[b.y]) return pricesA[a.x] < pricesA[b.x];
            return pricesA[a.x] + pricesB[a.y] < pricesA[b.x] + pricesB[b.y];
        };
        priority_queue<pii, vector<pii>, decltype(cmp)> q(cmp);
        vector<vector<int>> res;
        for (int i = 0; i < n; i ++) {
            q.emplace(i, m - 1);
        }
        while (k --) {
            auto [x, y] = q.top();
            q.pop();
            res.push_back({pricesA[x], pricesB[y]});
            if (y > 0) q.emplace(x, y - 1);
        }
        return res;
    }
};