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
廣告
資料結構
網路
關係資料庫管理系統 (RDBMS)
作業系統
Java
iOS
HTML
CSS
Android
Python
C語言程式設計
C++
C#
MongoDB
MySQL
Javascript
PHP