leetcode 373. 查找和最小的 K 对数字

128 阅读1分钟

c++

class Solution {
public:
    struct Node {
        int x, y, sum;
        Node (int x, int y, int sum) : x(x), y(y), sum(sum) {};
        bool operator< (const Node &a) const {
            return sum < a.sum;
        }
    };
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        priority_queue<Node> que;
        vector<vector<int>> ans;
        for (int i = 0; i < nums1.size(); i++) {
            for (int j = 0; j < nums2.size(); j++) {
                Node tmp(nums1[i], nums2[j], nums1[i] + nums2[j]); 
                if (que.size() < k || tmp.sum < que.top().sum) {
                    que.push(tmp);
                    if (que.size() > k) que.pop();
                } else break;
            }
        }
        while(que.size()) {
            Node t = que.top();
            que.pop();
            ans.push_back(vector({t.x, t.y}));
        }
        return ans;
    }
};

js

var kSmallestPairs = function(nums1, nums2, k) {
    var que = new MaxPriorityQueue({ priority : t => t[0] + t[1] });
    for (var x of nums1) {
        for (var y of nums2) {
            var tmp = [x, y];
            que.enqueue(tmp);
            if (que.size() > k) {
                if (que.dequeue().element === tmp) {
                    break;
                }
            }
        }
    }
    return que.toArray().map(i => i.element);
};