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

Arnab Chakraborty

更新于：2020-12-15

178 次浏览

开启您的职业生涯

完成课程获得认证

开始学习