2020年5月4日 星期一

Leetcode題解 Python:Find K Pairs with Smallest Sums

給兩個數列 nums1 nums2 ,要回傳 k 個總和由小排列的組合。參考

關於會作這一題, 是因為有一題更進階的在後頭,所以先把這個處理掉再進階。

之前我碰到 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