373.Find K Pairs with Smallest Sums

Posted on | Edited on | In , |

题目

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/

想法

一开始没有明确的想法，看了别人的思路之后明白…
用堆来实现，找最小值！
关键在于如何确定下一个可能的最小值！

  1. 先把(0,0), (1,0), …, (n,0)放进去
  2. 从heap中不断pop最小
  3. 最小(i,j)的下一个候选是(i, j+1), 关键
  4. 把下一个候选push进去就可以了

答案

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> res;
        if (nums1.empty() || nums2.empty() || k == 0)
            return res;
        auto comp = [nums1, nums2](pair<int, int> x, pair<int, int> y) {return nums1[x.first] + nums2[x.second] > nums1[y.first] + nums2[y.second];};
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(comp)> q(comp);
        for (int i = 0; i < nums1.size(); i++) {
            q.push(pair<int, int>(i, 0));
        }
        for (int i = 0; i < k; i++) {
            pair<int, int> p = q.top();
            q.pop();
            res.push_back(pair<int, int>(nums1[p.first], nums2[p.second]));
            if (p.first == nums1.size() - 1 && p.second == nums2.size() - 1)
                break;
            if (p.second < nums2.size() - 1)
                q.push(pair<int, int>(p.first, p.second+1));
        }
        return res;
    }
};

回顾

  1. 找准关系，如何能确定下一个和当前最小值的比较
  2. C++ priority_queue的使用