====== Find K Pairs with Smallest Sums ====== * [[ps:tutorial:fracturing search]] 의 가장 기본적인 형태이다. 이 개념을 몰라도 풀 수 있는 문제이지만, 여기에서는 fracturing search의 프레임워크로 설명하겠다. * [num1[i], num2[j]] 를 갖는 노드를 그냥 편의상 (i,j) 라고 표현하겠다 * 이미 정렬되어 있으므로 (0,0) 이 가장 값이 작고, (i,j) <= (i,j+1) 과 (i,j) <= (i+1,j) 을 모두 만족한다. * 모든 (i,j) 노드가 (i,j+1)과 (i+1,j) 을 자식으로 갖도록 트리를 만든다면, (i+1,j)와 (i,j+1)이 모두 (i+1,j+1)을 자식으로 갖게 되는 문제가 있다. * (0,j) 노드는 (1,j)와 (0,j+1)을 자식으로 갖고, i>0 인 i에 대해서는 (i,j)가 (i,j+1)만을 자식으로 갖도록 트리를 구성하면 중복 없이 모든 노드들을 힙 속성을 만족하는 트리로 만들 수 있다. * 이제 이 트리의 노드는 최대 2개 (=O(1)) 의 자식을 가지므로, fracturing search를 적용하면 O(klogk)에 답을 구할 수 있다 ===== 코드 ===== """Solution code for "LeetCode 373. Find K Pairs with Smallest Sums". - Problem link: https://leetcode.com/problems/find-k-pairs-with-smallest-sums - Solution link: http://www.teferi.net/ps/problems/leetcode/373 """ import heapq from typing import List class Solution: def kSmallestPairs( self, nums1: List[int], nums2: List[int], k: int ) -> List[List[int]]: heap = [(nums1[0] + nums2[0], 0, 0)] answer = [] for _ in range(k): _, ind1, ind2 = heapq.heappop(heap) answer.append([nums1[ind1], nums2[ind2]]) if ind2 == 0 and ind1 + 1 < len(nums1): heapq.heappush( heap, (nums1[ind1 + 1] + nums2[ind2], ind1 + 1, ind2) ) if ind2 + 1 < len(nums2): heapq.heappush( heap, (nums1[ind1] + nums2[ind2 + 1], ind1, ind2 + 1) ) return answer