Python程式:查詢K個最大和的數對


假設我們有兩個數字列表nums0和nums1,以及一個整數k。我們的目標是找到k個最大的數對之和,其中每個數對包含nums0中的一個整數和nums1中的另一個整數。需要返回所有數對的和。

因此,如果輸入類似於nums1 = [8, 6, 12],nums2 = [4, 6, 8],k = 2,則輸出將為38。我們有這些最大的數對[12, 8]和[12, 6]。

為了解決這個問題,我們將遵循以下步驟:

  • 如果k > len(nums0) * len(nums1)不為零,則

    • 返回0

  • pq := 新建一個最小堆

  • ans := 0

  • 對列表nums0和nums1進行排序

  • i, j := nums0的大小 - 1, nums1的大小 - 1

  • visited := 新建一個集合

  • 將-(nums0[i] + nums1[j]) , i, j壓入堆pq

  • 對於範圍0到k的i,執行以下操作:

    • sum, i, j := 從堆pq彈出

    • x := nums0[i - 1] + nums1[j]

    • 如果(i - 1, j)不在visited中,則

      • 將(i - 1, j)新增到visited

      • 將-x, i - 1, j壓入堆pq

    • y := nums0[i] + nums1[j - 1]

    • 如果(i, j - 1)不在visited中,則

      • 將(i, j - 1)新增到visited

      • 將-y, i, j - 1壓入堆pq

    • ans := ans + -sum

  • 返回ans

讓我們看看下面的實現,以便更好地理解:

Python

 線上演示

from heapq import heappush, heappop
class Solution:
   def solve(self, nums0, nums1, k):
      if k > len(nums0) * len(nums1):
         return 0
      pq = []
      ans = 0
      nums0.sort(), nums1.sort()
      i, j = len(nums0) − 1, len(nums1) − 1
      visited = set()
      heappush(pq, (−(nums0[i] + nums1[j]), i, j))
      for _ in range(k):
         sum, i, j = heappop(pq)
         x = nums0[i − 1] + nums1[j]
         if not (i − 1, j) in visited:
            visited.add((i − 1, j))
            heappush(pq, (−x, i − 1, j))
         y = nums0[i] + nums1[j − 1]
         if not (i, j − 1) in visited:
            visited.add((i, j − 1))
            heappush(pq, (−y, i, j − 1))
         ans += −sum
      return ans
ob = Solution()
print(ob.solve([8, 6, 12], [4, 6, 8], 2))

輸入

[8, 6, 12],[4, 6, 8],2

輸出

38

更新於:2020-12-15

178 次瀏覽

開啟您的職業生涯

完成課程獲得認證

開始學習
廣告
© . All rights reserved.