LeetCode Top Interview 150) Find K Pairs with Smallest Sums

Question

You are given two integer arrays nums1 and nums2 sorted in non-decreasing order and an integer k.

Define a pair (u, v) which consists of one element from the first array and one element from the second array.

Return the k pairs (u1, v1), (u2, v2), ..., (uk, vk) with the smallest sums.

Constraints:

• 1 <= nums1.length, nums2.length <= 10^5
• -10^9 <= nums1[i], nums2[i] <= 10^9
• nums1 and nums2 both are sorted in non-decreasing order.
• 1 <= k <= 10^4

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        

---

---

오름차순 정렬된 두 정수리스트 nums1, nums2에서 각각 원소를 하나씩 고릅니다. 이 때 합이 가장 작은 쌍을 k개 반환하는 문제입니다.
만약 k보다 순서쌍이 가능한만큼 순서쌍을 반환합니다.

제가 생각한 코드는 다음과 같습니다.

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        answer = []
        heap = [(n1+nums2[0],n1,nums2[0],0) for n1 in nums1[:k]]
        
        while k > 0 and heap:
            _, n1, n2, idx = heappop(heap)
            answer.append([n1, n2])

            if idx + 1 < len(nums2):
                heappush(heap, (n1+nums2[idx+1], n1, nums2[idx+1], idx+1))
                
            k -= 1
                
        return answer

• answer : 반환값을 저장할 변수입니다.
• heap : 합의 최솟값을 판단하기 위해 사용할 heap 리스트입니다.
  • 정수리스트가 이미 오름차순 정렬되어 있으므로 k번째 까지 nums1원소와 nums2의 첫번째 원소를 저장합니다.
  • heap은 원소의 첫번째 값을 기준으로 구조화됩니다.
  • (두 원소의 합, nums1 원소, nums2 원소, nums2 인덱스)를 저장합니다.
• k번 또는 heap이 빌 때 까지 while문을 반복합니다.
  • 반환리스트 answer에 현재 합의 최솟값인 [nums1 원소,nums2 원소]를 반환합니다.
  • 만약 현재 nums2의 다음 인덱스가 존재한다면 nums2의 다음 인덱스와 원소를 heap원소 구조에 맞게 추가합니다.
  • k번 순회를 위해 k을 1씩 빼줍니다.

---

문제 자체는 단순했지만 힙을 이용한 풀이가 잘 떠오르지 않았던 문제입니다.

오름차순 된 두 정수리스트니까 좀 더 쉽게 풀 수 있을거란 생각에 헛수고를 좀 해야했습니다. 단순히 인덱스로만 문제를 해결하려 했더니 조건이 너무 많아지고 nums1, nums2의 원소 합을 비교하는 방식까지 복잡해졌습니다.

이런 식으로 풀다가는 시간복잡도가 낮게 나와도 큰 의미가 없을듯해서 힙을 사용하는 풀이로 방향을 바꾸게 되었습니다.. ( 1시간반 헛수고요 ^~^)