[LeetCode] Find K Pairs with Smallest Sums 找和最小的K對數字
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.
Example 1:
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]
Example 2:
Given nums1 = [1,1,2], nums2 = [1,2,3], k = 2 Return: [1,1],[1,1] The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
Example 3:
Given nums1 = [1,2], nums2 = [3], k = 3 Return: [1,3],[2,3] All possible pairs are returned from the sequence: [1,3],[2,3]
Credits:
Special thanks to @elmirap and @StefanPochmann for adding this problem and creating all test cases.
這道題給了我們兩個陣列,讓我們從每個陣列中任意取出一個數字來組成不同的數字對,返回前K個和最小的數字對。那麼這道題有多種解法,我們首先來看brute force的解法,這種方法我們從0迴圈到陣列的個數和k之間的較小值,這樣做的好處是如果k遠小於陣列個數時,我們不需要計算所有的數字對,而是最多計算k*k個數字對,然後將其都儲存在res裡,這時候我們給res排序,用我們自定義的比較器,就是和的比較,然後把比k多出的數字對刪掉即可,參見程式碼如下:
解法一:
class Solution { public: vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) { vector<pair<int, int>> res; for (int i = 0; i < min((int)nums1.size(), k); ++i) { for (int j = 0; j < min((int)nums2.size(), k); ++j) { res.push_back({nums1[i], nums2[j]}); } } sort(res.begin(), res.end(), [](pair<int, int> &a, pair<int, int> &b){return a.first + a.second < b.first + b.second;}); if (res.size() > k) res.erase(res.begin() + k, res.end()); return res; } };
我們也可以使用multimap來做,思路是我們將陣列對之和作為key存入multimap中,利用其自動排序的機制,這樣我們就可以省去sort的步驟,最後把前k個存入res中即可:
解法二:
class Solution { public: vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) { vector<pair<int, int>> res; multimap<int, pair<int, int>> m; for (int i = 0; i < min((int)nums1.size(), k); ++i) { for (int j = 0; j < min((int)nums2.size(), k); ++j) { m.insert({nums1[i] + nums2[j], {nums1[i], nums2[j]}}); } } for (auto it = m.begin(); it != m.end(); ++it) { res.push_back(it->second); if (--k <= 0) return res; } return res; } };
下面這種方式用了priority_queue,也需要我們自定義比較器,整體思路和上面的沒有什麼區別:
解法三:
class Solution { public: vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) { vector<pair<int, int>> res; priority_queue<pair<int, int>, vector<pair<int, int>>, cmp> q; for (int i = 0; i < min((int)nums1.size(), k); ++i) { for (int j = 0; j < min((int)nums2.size(), k); ++j) { if (q.size() < k) { q.push({nums1[i], nums2[j]}); } else if (nums1[i] + nums2[j] < q.top().first + q.top().second) { q.push({nums1[i], nums2[j]}); q.pop(); } } } while (!q.empty()) { res.push_back(q.top()); q.pop(); } return res; } struct cmp { bool operator() (pair<int, int> &a, pair<int, int> &b) { return a.first + a.second < b.first + b.second; } }; };
下面這種方法比較另類,我們遍歷nums1陣列,對於nums1陣列中的每一個數字,我們並不需要遍歷nums2中所有的數字,實際上,對於nums1中的數字,我們只需要記錄nums2中下一個可能組成數字對的座標,這裡我們使用一個idx陣列,其中idx[i]表示的數字是nums1[i]將從nums2陣列上開始尋找的位置,因為 {nums1[i], nums2[i - 1]} 已經被加入到了結果res中,這種方法其實也是一種地毯式搜尋,但是並不需要遍歷完所有的組合,因為我們有idx陣列來進行剪枝。我們suppose需要進行k次迴圈,但是題目中沒有說我們一定能取出k對數字,而我們能取出的對兒數跟陣列nums1和nums2的長度有關,最多能取出二者的長度之積的對兒數,所以我們取其跟k之間的較小值為迴圈次數。我們定義idx陣列,長度為nums1的長度,初始化均為0。下面開始迴圈,在每次迴圈中,我們新建變數cur,記錄從nums1中取數的位置,初始化為0,使用變數sum來記錄一個當前最小的兩數之和,初始化為正無窮。然後開始遍歷陣列nums1,更新sum的條件有兩個,第一個是idx[i]上的數要小於nums2的長度,因為其是在nums2開始尋找的位置,當然不能越界,第二個條件的候選的兩個陣列 nums1[i] 和 nums2[idx[i]] 之和小於等於sum。同時滿足這兩個條件就可以更新sum了,同時更新cur為i,表示當前從nums1取出數字的位置。當遍歷nums1的for迴圈結束後,此時cur的位置就是要從nums1取出的數字的位置,根據idx[cur]從nums2中取出對應的陣列,形成數對兒存入結果res中,然後idx[cur]自增1,因為當前位置的數字已經用過了,下次遍歷直接從後面一個數字開始吧,這是本解法的設計精髓所在,一定要弄清楚idx陣列的意義,參見程式碼如下:
解法四:
class Solution { public: vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) { vector<pair<int, int>> res; int size = min(k, int(nums1.size() * nums2.size())); vector<int> idx(nums1.size(), 0); for (int t = 0; t < size; ++t) { int cur = 0, sum = INT_MAX; for (int i = 0; i < nums1.size(); ++i) { if (idx[i] < nums2.size() && sum >= nums1[i] + nums2[idx[i]]) { cur = i; sum = nums1[i] + nums2[idx[i]]; } } res.push_back({nums1[cur], nums2[idx[cur]]}); ++idx[cur]; } return res; } };
參考資料: