Leetcode 373. Find K Pairs with Smallest Sums

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/

Find K Pairs with Smallest Sums - LeetCode

문제

정수로 이루어진 오름차순 정렬된 배열 nums1과 nums2와 정수 k가 주어진다,

각 배열에서 추출된 정수로 (u , v)의 쌍을 정의하여 k 만큼의 쌍을 쌍의 합이 작은 순서대로 반환하라

 

코드

from heapq import *

class Solution:

    def kSmallestPairs(self, nums1: list[int], nums2: list[int], k: int) -> list[list[int]]:

        answer = []

        l_nums1, l_nums2, visited = len(nums1), len(nums2), set()

        if l_nums2 == 0 or l_nums1 == 0:

            return []

 

        h = [(nums1[0] + nums2[0], (0, 0))]

 

        for _ in range(min(k, l_nums1 * l_nums2)):

            s, (i, j) = heappop(h)

            answer += [[nums1[i], nums2[j]]]

 

            if i + 1 < l_nums1 and (i +1, j) not in visited:

                heappush(h, (nums1[i + 1] + nums2[j], (i + 1, j)))

                visited.add((i+1, j))

 

            if j + 1 < l_nums2 and (i , j + 1) not in visited:

                heappush(h, (nums1[i] + nums2[j + 1], (i,  j + 1)))

                visited.add((i, j + 1))

        return answer

---

우선순위 큐에 대해서 잘 몰랐었기에 다른 사람의 코드를 참조했다. 먼저 일반적으로 사람들이 말하는 우선순위큐는 python에서는 heapq를 이용해 쉽게 구현할 수 있다.먼저 heapq는 다음과 같은 규칙을 이용해 원소를 정리한다.

heap[k] <= heap[2*k+1], heap[k] <= heap[2*k+2]

따라서 0번째 인덱스에 위치한 값은 항상 최소라고 볼 수 있다. 따라서 tuple비교 연산의 경우 비교가 되는 튜플의 길이가 같은 경우 0번째 인덱스부터 비교해 값이 작은 순서대로 저장된다. 결국 이러한 점을 이용해 문제에서 제시하는 요건을 충족할 수 있음을 파악할 수 있다.

 

우선순위 큐를 사용하기에 min(k, len(nums1) * len(nums2))의 반복 횟수만큼 줄일 수 있다. 이러한 반복횟수 안에서 index가 배열을 넘지 않는 범위 안에서 반복적으로 heap에 값을 넣은 후 최소합을 가진 원소쌍을 answer 배열에 넣음으로써 답을 구현할 수 있다.

 

visited의 경우 i의 증가와 j의 증가가 따로 이루어짐에 따라서 (1,1) 또는 (2,2)가 중복해서 들어갈 수 있어 중복을 배제하기 위해 따로 관리하여 중복되지 않는 원소 쌍을 구할 수 있다.