關於會作這一題, 是因為有一題更進階的在後頭,所以先把這個處理掉再進階。
之前我碰到 heap 都不是了解,但是又很多人會用,所以這次我得學著用 heap 去解問題。
關於 heap , 是一種「堆積法」,有 maxheap 和 minheap ,可想成是一個二元樹,把根節設為最大或最小。
好處是可以快速找到 max 或是 min ,在要用到這類的計算排列時,heap 就很好用了。
python的 heapq 是以 minheap 為主的,跟一般排序的 maxheap 是不同的。
而 minheap,也是這題會用到的 heap ,要找最小值。
如果用 heapq.heappush() 相當於 append()、sort(),這樣在枚舉時也能加速。
這題要找組合,所以要保存總和由小到大的組合的 result 。
由於 nums1 跟 nums2 都已經排序過,所以我們可以得知每個的第[0]位相加會是最小總和。
接下來第二小的,不是 nums1 的索引+1,就是 nums2 的索引+1。
這樣子會陷入一個難題,到底重複出現搜尋過索引組合的要如何避免?最簡單的方法是拿索引組合用hash類型做檢查。
由於 heap 的排序若有二維是用第[0]位做判斷(跟list一樣),因此 heap 中元素的第[0]位是總和,這樣也能符合由小到大的要求。
而後面幾位擺上索引組合,因為元素可能是相同數值,若只計算加總值可能會漏掉該出現的重複。
(這題的解法,已針對後面進階題做改良。)
Python
class Solution: def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]: result = [] if not nums1 or not nums2: return result nums = [nums1, nums2] heap = [(sum(each[0] for each in nums) , 0, 0)] memo = {(0,0)} n = len(nums) while k > 0 and heap: _, *idxs = heapq.heappop(heap) result.append([nums[j][idxs[j]] for j in range(n)]) k -= 1 for i in range(len(idxs)): idxs[i] += 1 if idxs[i] < len(nums[i]) and tuple(idxs) not in memo: memo.add(tuple(idxs)) heapq.heappush(heap, [sum([nums[j][idxs[j]] for j in range(n)]), *idxs]) idxs[i] -= 1 return result