502. IPO

Python

Max Heap

• 官方解法
• 在每個步驟，有兩個子問題
  • 找到下一個可承接的案子
    • 下一個所需預算最少的案子
    • 先將預算以小到大排序過
  • 找到當前預算可承擔的狀況下，能夠接下的最有profit的專案
    • 使用max heap維護剩餘的案子

from heapq import heappush, heappop


class Solution:
    def findMaximizedCapital(self, k: int, w: int, profits: List[int], capitals: List[int]) -> int:
        fund = w

        projects = list(zip(capitals, profits))
        projects.sort()

        max_heap = []
        ptr = 0

        for i in range(k):
            while ptr < len(projects) and projects[ptr][0] <= fund:
                profit = projects[ptr][1]
                heappush(max_heap, -profit) # Add new available project
                ptr += 1

            if not max_heap:
                break

            # Choose the best project and remove it
            fund += -heappop(max_heap)

        return fund