heap 자료구조 관련 leetcode 문제풀이

Heap이란?

• heap은 우선 순위 큐를 구현하는 자료구조로 이진 트리를 기반으로 구현된다. 
• 부모 노드가 자식 노드보다 큰 경우 최대 힙, 부모 노드가 자식 노드보다 작은 경우 최소 힙이라고 한다. 
• 힙의 삽입 및 삭제 연산은 O(logn)의 시간 복잡도로 이루어진다. 

 

파이썬에 내부에서 사용 방식

• heapq 모듈을 이용하여 사용할 수 있다. 
• 삽입 (heappush), 삭제 (heappop), 특정 배열을 힙 형태로 만들어주는 함수 (heapify) 지원
• heappop을 사용하면 배열 내의 가장 작은 값이 나오지만, 값이 삭제된다. 만약 삭제되길 원하지 않는다면 배열의 첫번째 원소에 접근하는 방식으로도 사용 가능하다. 
• 기본적으로 최소힙만 지원하기 때문에 최대힙을 구현하기 위해서는 들어가는 값에 -를 곱해주어야 한다. 

import heapq

my_heap = []

heapq.heappush(my_heap, 1)
heapq.heappush(my_heap, 3)
heapq.heappush(my_heap, 2)

small = heapq.heappop(my_heap) # small = 1

small2 = my_heap[0] # small2 = 2, 최솟값에 배열 형태로도 접근가능

 

 

관련 문제

1. Leetcode 502번 IPO

https://leetcode.com/problems/ipo/description/?envType=study-plan-v2&envId=top-interview-150

 

특정 프로젝트를 진행할 때, w원 이상의 자본이 있을 경우에만 진행이 가능하다. 

프로젝트를 진행하면 profit 만큼의 자본을 획득하게 된다고 할 때, k개의 프로젝트를 진행해서 가장 많은 자본을 획득하는 방법은?

 

class Solution(object):
    def findMaximizedCapital(self, k, w, profits, capital):
        """
        :type k: int
        :type w: int
        :type profits: List[int]
        :type capital: List[int]
        :rtype: int
        """

        done = [False for i in range(len(profits))]
        pos_pro = []
        impos_pro = []

        for index, (profit, cap) in enumerate(zip(profits, capital)):
            if cap <= w:
                heapq.heappush(pos_pro, (-profit, -cap, index))
            else:
                heapq.heappush(impos_pro, (cap, -profit, index))
        

        for project_num in range(k):
            # get project
            while True:
                if len(pos_pro) == 0:
                    break
                else:
                    project = heappop(pos_pro)
                    profit, cap, index = -project[0], -project[1], project[2]
                    
                    if done[index] is False:
                        done[index] = True
                        w += profit
                        break
            
                    

            # get possible projects
            while True:
                if len(impos_pro) == 0:
                    break
                else:
                    project = impos_pro[0]
                    cap, profit, index = project[0], -project[1], project[2]
                    if cap <= w:
                        heapq.heappush(pos_pro, (-profit, -cap, index))
                        heapq.heappop(impos_pro)
                    else:
                        break


        return w

 

프로젝트 정보를 저장하는 힙을 사용하여 가장 많은 자본을 가질 수 있는 프로젝트를 heappop을 통해서 얻는다. 

하지만, 그 과정에서 내가 가진 자본 내에서 진행할 수 있는지 체크 + 프로젝트가 이미 진행되었는지 여부 체크가 필요하다. 

 

힙과 프로젝트의 완료 여부를 저장하는 인덱스를 함께 사용하여야 하는 문제

 

2. Leetcode 373번 Find K Pairs with Smallest Sums

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/?envType=study-plan-v2&envId=top-interview-150

 

2개의 정렬된 배열이 주어진다. 

각 배열에서 1가지의 원소씩 뽑아 길이가 2인 배열을 만든다. 

뽑힌 2원소의 합이 가장 작은 순서대로 k개를 뽑으시오. 

 

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """

        min_heap = []
        answers = []
        visited = set()
        heapq.heappush(min_heap, (nums1[0] + nums2[0], 0, 0))
        visited.add((0, 0))
        for i in range(k):
            # print(min_heap)
            answer = heapq.heappop(min_heap)
            i1, i2 = answer[1], answer[2]
            num1, num2 = nums1[i1], nums2[i2]
            answers.append([num1, num2])

            if i1 + 1 < len(nums1) and (i1 + 1, i2) not in visited:
                ni1 = i1 + 1
                ni2 = i2
                heapq.heappush(min_heap, (nums1[ni1] + nums2[ni2], ni1, ni2))
                visited.add((ni1, ni2))
            
            if i2 + 1 < len(nums2) and (i1, i2 + 1) not in visited:
                ni1 = i1
                ni2 = i2 + 1
                heapq.heappush(min_heap, (nums1[ni1] + nums2[ni2], ni1, ni2))
                visited.add((ni1, ni2))
        
        return answers

 

각 배열의 (i, j) 인덱스에 있는 값의 합이 지금까지 본 값중에서 가장 최솟값이었다면

그 다음으로 작은 것은

1. 각 배열에서 (i +1, j) 인덱스에 있는 값의 합

2. 각 배열에서 (i, j + 1) 인덱스에 있는 값의 합

둘 중에 하나가 된다. 

 

그렇기 때문에 각 값을 힙에 넣어 가장 합이 작은 값을 얻을 수 있도록 한다. 

(BFS인데 큐를 우선순위 큐를 쓰는 느낌)