[Leetcode] 502 IPO

https://leetcode.com/problems/ipo/

leetcode에서 IPO 준비를 위한 자본금 확보를 위해 프로젝트를 진행하는데, 프로젝트와 각각 필요한 자본이 주어지고 진행할 수 있는 프로젝트의 개수가 주어질 때 최대의 자본을 확보하라는 문제이다.

코드 1 - 실패

from collections import defaultdict
from heapq import heappush, heappop

class Solution:
    def findMaximizedCapital(self, k: int, w: int, profits: List[int], capital: List[int]) -> int:
        pc = defaultdict(list)

        for idx, c in enumerate(capital):
            heappush(pc[c], -profits[idx])

        for i in range(k):
            cap = w
            maxPro = 0
            maxCap = cap
            while cap >= 0:
                if pc[cap]:
                    if maxPro < -min(pc[cap]):
                        maxPro = -min(pc[cap])
                        maxCap = cap
                cap -= 1

            # print(maxCap)
            
            if cap == -1 and maxPro == 0:
                break

            # print(i, w, maxCap)
            w += -heappop(pc[maxCap])
            
        
        return w

필요 자본을 key로하는 딕셔너리를 만들고 내부에는 수익을 기준으로 하는 힙을 구성했다.
딕셔너리를 탐색하면서 현재 자본 이하로 진행할 수 있는 프로젝트 중 어떤 프로젝트가 최고인지 찾아서 더하고 하는 방식으로 진행했는데 당연히도 실패했다.

코드 2

그래서 전체를 힙으로 넣고 탐색해보기로 했다.
minHeap과 maxHeap을 두고 최초에는 maxHeap을 탐색하면서 지금 자본금으로 할 수 있는 프로젝트를 찾으면 더하고, 진행할 수 없으면 minHeap에 보관하고 하는 방식으로 해보려고 한다.

from heapq import heappop, heappush

class Solution:
    def findMaximizedCapital(self, k: int, w: int, profits: List[int], capital: List[int]) -> int:
        maxHeap = []
        minHeap = []

        for idx, c in enumerate(capital):
            heappush(maxHeap, (-profits[idx], c))

        for _ in range(k):
            while minHeap and minHeap[0][0] <= w:
                c, p = heappop(minHeap)
                heappush(maxHeap,(p, c))

            while maxHeap and maxHeap[0][1] > w:
                p, c = heappop(maxHeap)
                heappush(minHeap, (c, p))
            
            if not maxHeap:
                break
            
            pro, cap = heappop(maxHeap)

            w -= pro
        return w

현재 자본보다 큰 프로젝트는 minHeap에 저장했다가 매번 시작할 때 minHeap에서 지금 자본보다 작은 프로젝트 다시 꺼내서 maxHeap으로 옮겼다.
그리고 다시 maxHeap에서는 수익 큰 프로젝트부터 현재 자본과 필요자본을 비교하면서 꺼내서 minHeap으로 저장하다가 할 수 있는 프로젝트 나오면 프로젝트 진행하고 수익 얻고 하는 방식으로 진행했고 통과할 수 있었다.