373. Find K Pairs with Smallest Sums

📘 題目敘述

給定兩個已經按照非遞減順序排序好的整數陣列 nums1 和 nums2，以及一個整數 k。

定義一組 pair (u, v)：

• u 來自 nums1
• v 來自 nums2

請回傳總和最小的前 k 組 pair。

條件限制

• 1 <= nums1.length, nums2.length <= 10^5
• -10^9 <= nums1[i], nums2[i] <= 10^9
• nums1 和 nums2 都已經按照非遞減順序排列
• 1 <= k <= 10^4
• k <= nums1.length * nums2.length

🧠 解題思路

我這題的想法是把整個配對空間想成一個二維表格。

如果我固定：

• 第 i 列代表 nums1[i]
• 第 j 行代表 nums2[j]

那每個位置 (i, j) 就對應到一組 pair：

• (nums1[i], nums2[j])

而且因為兩個陣列都已經排序好了，所以：

• 往下走，pair sum 只會不變或變大
• 往右走，pair sum 也只會不變或變大

所以我可以從最小的 (0, 0) 開始，用 min heap 每次取出目前最小的 pair。當我取出 (i, j) 之後，再把它的「下方 (i + 1, j)」和「右方 (i, j + 1)」當成新的候選人丟進 heap。

這樣就像是在一個遞增的表格中做 best-first search，逐步找出前 k 小的 pair。

另外因為同一個位置可能會從不同路徑被推進 heap，所以我再用一個 set 來避免重複加入相同座標。

所有變數

• n1：nums1 的長度
• n2：nums2 的長度
• check：記錄哪些座標 (i, j) 已經進過 heap，避免重複加入
• pq：min heap，裡面放 {pair sum, i, j}，用來每次取出目前總和最小的候選 pair
• ans：最後要回傳的前 k 組 pair
• x：從 heap 取出的目前最小候選資料
• i：目前取出的 pair 在 nums1 中的索引
• j：目前取出的 pair 在 nums2 中的索引

🪜 主要流程步驟

1. 先建立最小候選起點 (0, 0)

因為兩個陣列都已經排序好，所以整個表格中最小的 pair 一定是：

• nums1[0] + nums2[0]

所以我先把 (0, 0) 放進 min heap：

pq.push({nums1[0] + nums2[0], 0, 0});

這代表我的搜尋會從整張表格左上角，也就是最小的起點開始。

2. 每次從 heap 取出目前總和最小的 pair

我接著一直做，直到已經取出 k 組 pair 為止。

每次先從 heap 取出最小的候選：

auto x = pq.top();
pq.pop();
int i = x[1];
int j = x[2];

因為 pq 裡放的是 {總和, i, j}，而且用的是 greater<vector<int>>，所以 heap 頂端就是目前總和最小的那組 pair。

取出後，我就把對應的實際數值加入答案

3. 取出 (i, j) 之後，往下擴展 (i + 1, j)

如果目前位置是 (i, j)，那麼它下面那格 (i + 1, j) 也是一個合理的下一個候選。

因為 nums1 已經排序好，所以從 (i, j) 往下走時，pair sum 不會變小。這表示它有可能成為之後的最小值，所以要把它丟進 heap。

這裡除了檢查邊界，也要確認這個座標之前沒有進過 heap，避免重複。

4. 取出 (i, j) 之後，往右擴展 (i, j + 1)

同樣地，右邊那格 (i, j + 1) 也是下一個可能的候選。

因為 nums2 也已經排序好，所以往右走時，pair sum 也不會變小。所以它也應該加入 heap：

這樣一來，每次從當前最小位置往右和往下擴展，就能逐漸把整張表格中比較小的區域探索出來。

5. 用 check 避免同一個座標重複進 heap

這題如果不做去重，某些位置可能會被重複加入。

例如某個 (i, j) 可能同時從：

• (i - 1, j) 往下走過來
• (i, j - 1) 往右走過來

如果不擋掉，這個位置就會進 heap 兩次。所以我用：

set<pair<int, int>> check;

來記錄哪些座標已經進過 heap。

只要某個位置已經被加入過，就不再重複 push。這樣才能保證每個位置最多只進 heap 一次。

6. 重複取出 k 次後回傳答案

每次從 heap 取出一組最小 pair，就把 k 減一：

k--;

當 k 變成 0，表示前 k 小的 pair 都已經收集完了，這時直接回傳 ans 就可以。

⏱️ 複雜度

設實際加入 heap 的座標數量為 K，其中 K <= k + 1 的同階量級。

• 時間複雜度：O(k log k) 每次取出一組 pair，都最多新增兩個新候選，而 heap 與 set 操作都是對數時間。
• 空間複雜度：O(k) pq、check 和 ans 都最多只會存和前 k 組答案同階數量的資料。

💻 程式碼實作

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        int n1 = nums1.size();
        int n2 = nums2.size();
        set<pair<int, int>> check;
        priority_queue<vector<int>, vector<vector<int>>, greater<vector<int>>> pq;
        pq.push({nums1[0] + nums2[0], 0, 0});
        vector<vector<int>> ans;

        while (k > 0) {
            auto x = pq.top();
            pq.pop();
            int i = x[1];
            int j = x[2];
            ans.push_back({nums1[i], nums2[j]});
            if (i + 1 < n1 && !check.count({i + 1, j})) {
                pq.push({nums1[i + 1] + nums2[j], i + 1, j});
                check.insert({i + 1, j});
            }
            if (j + 1 < n2 && !check.count({i, j + 1})) {
                pq.push({nums1[i] + nums2[j + 1], i, j + 1});
                check.insert({i, j + 1});
            }
            k--;
        }
        return ans;
    }
};

🔗 題目連結

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