373. Find K Pairs with Smallest Sums

harry xie
2024-10-05 00:22:38

解題程式碼

var kSmallestPairs = function (nums1, nums2, k) {
  const res = [];
  const minHeap = new MinPriorityQueue({ compare: (a, b) => a[0] - b[0] });
  const visited = new Set();
  minHeap.enqueue([nums1[0] + nums2[0], 0, 0]);

  while (minHeap.size() && res.length < k) {
    const [sum, i, j] = minHeap.dequeue();

    if (visited.has(`${i}-${j}`)) continue;
    visited.add(`${i}-${j}`);
    res.push([nums1[i], nums2[j]]);

    if (i + 1 < nums1.length) minHeap.enqueue([nums1[i + 1] + nums2[j], i + 1, j]);
    if (j + 1 < nums2.length) minHeap.enqueue([nums1[i] + nums2[j + 1], i, j + 1]);
  }

  return res;
};

解題思路演算法

要找最小值的資料結構，可以想到 Heap

而 nums1[0] + nums[0] 的總和一定是最小的 pair，所以一開始可以建立 Heap 並將兩個索引和總和 [sum, nums1 index, nums2 index] 加入 Heap

接著因為它在當前的 Heap 最小，所以可以加到題目要的最終結果陣列 res，然後從表格 [0, 0] 的位置開始往右邊和下面展開，將兩個值加入 Heap

來到第二次迴圈，一樣找在當前的 Heap 最小的 sum，並從該 sum 索引開始往右邊和下面展開，將兩個值加入 Heap若有重複走過的，則不需加入，所以需要一個 hashSet 去儲存

重複以上過程，直到找到 k 個 pairs

如圖，k = 4，一開始找到紅色區塊的最小值 [1, 2]，接著延伸到綠色區塊，從填色區塊中找到最小值 [1, 4]，延伸到黃色區塊

再從填色區塊中找到最小值 [4, 2]，延伸到藍色區塊，[4, 4] 因為被延伸過所以不重複走，最後找到 [1, 6]，已經找到 4 個 pairs 所以就不再查找

解法的時間空間複雜度

時間複雜度: O(k * log k)
空間複雜度: O(k)

參考資料

補充官方題解優先隊列

373. Find K Pairs with Smallest Sums 查找和最小的 k 对数字LeetCode 单题讲解系列

373. Find K Pairs with Smallest Sums 0466