373. Find K Pairs with Smallest Sums

題目網址：https://leetcode.cn/problems/find-k-pairs-with-smallest-sums/

題意：給兩個升序排列的 array nums1 和 nums2。

定義 pair (u, v), 其中 u 來自 nums1, v 來自 nums2

求和最小的 k 個 pair

Solution 1：

想法：利用 Heap, 解法同 378. Kth Smallest Element in a Sorted Matrix, pair 分別存 nums1 第 i 個元素、nums2 第 j 個元素 (只存 index, 不存 val), 然後用 minHeap 取最小的 sum(nums1[i], nums2[j])

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        // 重要: 將 nums1, nums2 傳入 cmp 中
        const auto cmp = [&nums1, &nums2](const pair<int, int>& a, const pair<int, int>& b) {
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        };

        const int m = nums1.size();
        const int n = nums2.size();
        vector<vector<int>> res;

        // 重要: 自定義 priority_queue 的 cmp 函數, 請牢記 heap 預設為 !cmp
        // 故 minHeap 要用 greater<T>, 且 cmp 裡面是判斷第一個 > 第二個
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pq(cmp);

        // 將第一行元素加到 minHeap 中, k 的個數有可能小於 m
        for (int i = 0; i < min(k, m); i++) {
            pq.emplace(i, 0);
        }

        while (k-- > 0 && !pq.empty()) {
            const auto [i, j] = pq.top();
            pq.pop();
            res.emplace_back(vector<int>{nums1[i], nums2[j]});

            if (j + 1 < n) {
                pq.emplace(i, j + 1);
            }
        }

        return res;
    }
};

• time：$O(k \cdot log(k))$
  • $O(k)$：while loop, 其中 worse case 為 k = mn
  • $O(log(k))$：insertion of heap(因為 heap 是 balanced tree, swap() 最多為樹高 $log(k)$ 次)
• space：$O(k)$ ➔ minHeap 最多 k 個元素

Solution 2：

想法：利用 Binary Search, 解法同 378. Kth Smallest Element in a Sorted Matrix

class Solution {
private:
    long count(vector<int>& nums1, vector<int>& nums2, int& mid){
        // 只需 O(m+n) time, 因為 nums1 遞增, 若 i 變大, 則 j 變小
        // 假設 i = 0 做完, j = 10
        // 則 i = 1 時, j 從 10 繼續(pos 不用重設成 nums2.size() - 1)
        long cnt = 0;
        int j = nums2.size() - 1;

        for (int i = 0; i < nums1.size(); ++i) {
            while (j >= 0 && nums1[i] + nums2[j] > mid) {
                --j;
            }
            cnt += (j + 1);
        }

        return cnt;
    }

public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        long left = nums1[0] + nums2[0];
        long right = nums1.back() + nums2.back();

        while (left < right) {
            const int mid = left + (right - left) / 2;

            if (count(nums1, nums2, mid) < k) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }

        const int pairSum = left;
        vector<vector<int>>res1, res2;

        for (int i = 0; i < nums1.size(); ++i) {
            for (int j = 0; j < nums2.size() && nums1[i] + nums2[j] <= pairSum; ++j) {
                // sum < pairSum 的 pair 一定會取
                if (nums1[i] + nums2[j] < pairSum) {
                    res1.emplace_back(vector<int>{nums1[i], nums2[j]});
                } else {
                    // sum == pairSum 的 pair 不一定會取(有可能很多個, 導致最後超過 k 個)
                    // 故用另一個 vector 存
                    res2.emplace_back(vector<int>{nums1[i], nums2[j]});
                }
            }
        }

        // 從 res2 中取 k - res1.size() 個
        // 也有可能 res2 沒那麼多個給你取(k取太大, 沒那麼多pair), 故要取 min
        const int remain = min(k - res1.size(), res2.size());
        for (int i = 0; i < remain; ++i) {
            res1.emplace_back(res2[i]);
        }

        return res1;
    }
};

• time：$O((m+n) \cdot log(r-l) + m \cdot n)$ ➔ m 為 nums1.size(), n 為 nums2.size()
  • $O(log(r-l))$：Binary Search 次數
  • $O(m+n)$：count() 所需時間
  • $O(m \cdot n)$：for loop, 有用 nums1[i] + nums2[j] ≤ pairSum 做 pruning, 實際遠小於 mn 次
• space：$O(1)$ ➔ 扣除要返回的 array, 只需常數空間