Find K Pairs with Smallest Sums

문제

오름차순으로 정렬된 두 개의 정수 배열 nums1 및 nums2와 정수 k가 제공됩니다. 첫 번째 배열의 한 요소와 두 번째 배열의 한 요소로 구성된 쌍 (u, v)를 정의합니다. 합이 가장 작은 k쌍 (u1, v1), (u2, v2), ..., (uk, vk)를 반환합니다.

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]
Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]
Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

다른 사람의 풀이

from heapq import heappush, heappop

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        m = len(nums1) # nums1 배열의 길이
        n = len(nums2) # nums2 배열의 길이
        ans = [] # 결과 배열

        # nums1과 nums2의 인덱스를 넣을 자료구조, 중복 제거
        visited = set()

        # 두 배열 요소의 합과 두 배열의 각 인덱스를 튜플로 묶어서 minHeap 초기화
        minHeap = [(nums1[0] + nums2[0], (0, 0))]
        visited.add((0, 0)) # 인덱스 삽입

        # k가 0보다 크면서 minHeap 요소가 존재할 경우 순회
        while k > 0 and minHeap:

          # minHeap에서 가장 작은 요소의 합과 각 배열의 인덱스를 변수에 저장
          val, (i, j) = heappop(minHeap)
          # 결과의 각 배열 요소를 추가
          ans.append([nums1[i], nums2[j]])

          # i + 1이 m보다 작으면서 (i+1, j)이 visited 세트 자료형에 없을 경우
          # nums1 배열에서 다음 인덱스
          if i+1 < m and (i+1, j) not in visited:
            heappush(minHeap, (nums1[i+1] + nums2[j], (i+1, j)))
            visited.add((i+1, j))

          # j + 1이 n보다 작으면서 (i, j+1)이 visited 세트 자료형에 없을 경우
          # nums2 배열에서 다음 인덱스
          if j+1 < n and (i, j+1) not in visited:
            heappush(minHeap, (nums1[i] + nums2[j+1], (i, j+1)))
            visited.add((i, j+1))

          # k번 반복
          k -= 1

        return ans

최소힙을 사용한 풀이로 초기에 두 배열에서 가장 작은 요소들의 합과 배열의 각 인덱스를 튜플로 묶어 초기화한 뒤에 minHeap에 요소가 존재하지 않을 때까지 k번 반복하면서 가장 작은 조합 순으로 ans 배열에 추가하고 다음번에는 그 다음으로 작은 합 순으로 ans 배열에 k번 추가하게 됩니다.