Find K Pairs with Smallest Sums

• [작성중] 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