[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;
    }
};

 

參考資料：

https://discuss.leetcode.com/topic/50429/c-solution

https://discuss.leetcode.com/topic/50459/c-idea-of-using-multimap

https://discuss.leetcode.com/topic/50422/naive-accepted-solution-c

https://discuss.leetcode.com/topic/50421/c-priority_queue-solution

 

LeetCode All in One 題目講解匯總(持續更新中...)