leetcode Find K Pairs with Smallest Sums

不得不说,用惯了python之后 c++感到非常非常的陌生。打个最大堆居然都会遇到各种bug

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) …(uk,vk) with the smallest sums.
看了测例就很好理解

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

主要的做题思路就是用一个大小为k 的最大堆,先把它填满,然后用num1[i] + num2[j] 和top 比较,如果小就替换。

class Solution {
public:
    struct com {
        bool operator() (pair<int, int> &a, pair<int, int> &b) {
            return (a.first + a.second) < (b.first + b.second);
        }
    };
    vector<pair<int, int>> kSmallestPairs(vector<int>& n1, vector<int>& n2, int k) {
        priority_queue< pair<int, int>, vector<pair<int,int> > , com> pq;
        for (int i = 0; i < min((int)n1.size(), k); i++) {
            for (int j = 0; j < min((int)n2.size(), k); j++) {
                if (pq.size() < k) {
                    pq.push(make_pair(n1[i], n2[j]));
                } else if ((n1[i] + n2[j]) < (pq.top().first + pq.top().second)) {
                    pq.pop();
                    pq.push(make_pair(n1[i], n2[j]));
                }
            }
        }
        vector<pair<int,int> > res(pq.size());
        int i = pq.size() - 1;
        while(!pq.empty()) {
            res[i--] = pq.top();
            pq.pop();
        }
        return res;
    }
};

代码其实不长,看下应该就很好理解了。
主要mark 下函数重载,priority_queue 的知识。
因为priority_queue 是按照数据的优先级来进行排序的。每次push 和 pop 都会按照按照数据大小从大到小进行排序。

priority_queue<int, vector<int>, greater<int>, 自定义函数 >
我们可以通过自定义函数来规定数据的优先级

一开始我是想用重载 < 实现的,结果发现居然不行。原因是std 里面自带pair 的<实现,所以不会发现和调用我自己定义的重载,改成operate() 就可以了。 
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章