LeetCode 373. Find K Pairs with Smallest Sums (python)

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.

 

 

 

풀이

 

2개의 배열이 주어지고 각각 오름차순 되어있는 배열에서 인자를 뽑아 2개의 조합만들어 질 때  가장 작은 합을 갖는 k개의 쌍을 찾는 문제이다.

 

생각한 방법

1. 모든 2개의 조합을 만들어  heap을 사용해 (nums1+nums2,nums1,nums2) 형식으로 삽입하여, 합이 가장 작은 k개를 뽑는다.

 

 

2. 문제에서 원하는 쌍은 k개이니 두 개의 배열을 오름차순 정렬되어있는걸 활용하여  범위를 k까지 지정하여 조합을 만든다.

import heapq
class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        if not nums1 or not nums2:
            return
        heap = []
        for i in range(min(k, len(nums1))):
            for j in range(min(k, len(nums2))):
                heapq.heappush(heap,(nums1[i]+ nums2[j],[nums1[i],nums2[j]]))
        result =[]
        for __ in range(min(k, len(heap))):
            result.append(heapq.heappop(heap))
        
        return [i[1] for i in result]

 

 

실패 ) 메모리초과 

 

메모리값을 줄 일 수 있는 방법이 생각이 안나서 같은 조의 팀원의 코드를 가져와봤다.

 

result가 추가되는것을 제한한다. heap에 k개 까지만 저장한 뒤 k개가 넘는다면 최대힙을 통해서 가장 큰 값을 뺀다.

 

 

다른 사람의 코드 

import heapq
class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        if not nums1 or not nums2:
            return
        heap = []
        for i in range(min(k, len(nums1))):
            for j in range(min(k, len(nums2))):
                if len(heap)<k:
                    heapq.heappush(heap,(-(nums1[i]+ nums2[j]),[nums1[i],nums2[j]]))
                else:
                    if heap[0][0] < -(nums1[i] + nums2[j]):
                        heapq.heappop(heap)
                        heapq.heappush(heap,(-(nums1[i]+ nums2[j]),[nums1[i],nums2[j]]))
                    else:
                        break
        return [i[1] for i in heap]

 

메모리 공간을 위해서 heap에는 k개의 원소만 유지하며, k개의 이상의 인자가 더해진다면 기존에 있던 값 (최소 힙트리의 root값)과 비교하여 작은지 확인 한 후 기존 값보다 작다면 pop을 진행하고 다시 새로운 인자를 넣고,

아닐경우 멈춘다.  

 

다른사람의 코드

 

 

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        flag = (n := len(nums1)) > (m := len(nums2)) # 원소 개수 체크
        if flag:
            n, m, nums1, nums2 = m, n, nums2, nums1	# 원소 크기에 따라 heappush에서 앞에 들어갈지 뒤에 들어가리 정해준다,
            
        heap = []
        for i in range(min(k, n)):
            heappush(heap, (nums1[i] + nums2[0], i, 0)) # 
        
        ans = []
        
        while heap and len(ans) < k:
            s, i, j = heappop(heap)
            ans.append([nums1[i], nums2[j]] if not flag else [nums2[j], nums1[i]]) # 둘중 개수가 큰걸 앞으로 놓는다,
            
            if j + 1 < min(k, m):
                heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
                #개수가 적은쪽의 인덱스를 늘려가며 ans에 넣을 인자를 heap에 먼저 넣어줌 
            
        return ans