Leetcode 373. Find K Pairs with Smallest Sums

结课之后其实也陆陆续续写了一些题，到现在一共有60多了，不过一直没有抽出时间来写题解；所以趁放假还在学校的这段时间，抽一些当初被坑过的题目来写写 ^_^

题目大意

给出两个有序序列nums1和nums2，找出 k 个和最小的对(num1, num2)，使得num1来自nums1，num2来自nums2。

注意给出的 k 可能大于所有值对的数目。

解题思路

容易发现，最终得到的 k 个对必定有如下形式：

// 以a指代nums1，以b指代nums2
(a[0], b[0]), (a[0], b[1]), ..., (a[0], b[n0-1])
(a[1], b[0]), (a[1], b[1]), ..., (a[1], b[n1-1])
...
(a[i], b[0]), (a[i], b[1]), ..., (a[i], b[ni-1]) 

并且有n0+n1+...+ni==k，且只能保证同一行是递增的，同一列是递增的，但不同行列之间的大小无法确定。也即是说，我们没有办法判断每一行究竟要取多少个对，不能依行的次序进行求解。

这个时候考虑优先队列，显然有 i < k，那么我们只要将前 k 行的第一个元素都压入队列，形成小顶堆；每次取堆顶，并生成堆顶元素的下一个元素压入队列中，可以从小到大取到 k 个对并且保证不会错过任何最小对。

复杂度分析

优先队列的一次插入操作所需要的时间复杂度为 O(logn)，那么初始化优先队列的时间复杂度为 O(log1) + O(log2) + ...+ O(log(k-1))；而后依次取 k 个对，每次都需要进行一次删除操作和一次插入操作，时间复杂度为 2k*O(logk)，总时间复杂度为O(klogk)。

源代码

struct Pair {
    int i1, i2, a, b, sum;
    Pair(int i1, int i2, int a, int b) :
        i1(i1), i2(i2), a(a), b(b) {
            sum = a + b;
        }
};

class Comp {
public:
    bool operator()(const Pair &a, const Pair &b) {
        return a.sum >= b.sum;
    }
};

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int> > result;
        if (nums1.empty() || nums2.empty()) {
            return result;
        }

        priority_queue<Pair, vector<Pair>, Comp> pq;
        for (int i = 0; i < k && i < nums1.size(); i++) {
            pq.push(Pair(i, 0, nums1[i], nums2[0]));
        }

        int counter = 0;
        int index1, index2, a, b;

        while (counter < k && !pq.empty()) {
            counter++;
            index1 = pq.top().i1;
            index2 = pq.top().i2;
            a = pq.top().a;
            b = pq.top().b;
            result.push_back(pair<int, int>(a, b));
            pq.pop();

            if (index2 + 1 < nums2.size()) {
                pq.push(Pair(index1, index2 + 1, nums1[index1], nums2[index2 + 1]));
            }
        }

        return result;
    }
};