Leetcode 502. IPO

Problem

Leetcode 502 - IPO

Example:

Input: k = 2, w = 0, profits = [1,2,3], capital = [0,1,1]
Output: 4
Explanation: Since your initial capital is 0, you can only start the project indexed 0.
After finishing it you will obtain profit 1 and your capital becomes 1.
With capital 1, you can either start the project indexed 1 or the project indexed 2.
Since you can choose at most 2 projects, you need to finish the project indexed 2 to get the maximum capital.
Therefore, output the final maximized capital, which is 0 + 1 + 3 = 4.

Input: k = 3, w = 0, profits = [1,2,3], capital = [0,1,2]
Output: 6

Approach

We can solve this problem using 2 heaps, one for tracking the minimum capital, and one for tracking maximum profit.

minCap

• tracks the minimum capital paired with maximum profit
• if minCap[0][0] is <= w and minCap exists
  • pop from minCap then push profit into maxProfit

maxProfit

• tracks the maximum profit at that capital

Here is the Python code for the solution:

class Solution:
    def findMaximizedCapital(self, k: int, w: int, profits: List[int], capital: List[int]) -> int:
        maxProfit = []
        minCap = [(c, p) for c, p in zip(capital, profits)]
        heapq.heapify(minCap)

        for _ in range(k):
            while minCap and minCap[0][0] <= w:
                c, p = heapq.heappop(minCap)
                heapq.heappush(maxProfit, -p)
            
            if not maxProfit:
                break
            
            w += -heapq.heappop(maxProfit) 
        
        return w
    

Time Complexity and Space Complexity

Time Complexity: $O(n log n)$

Space Complexity: $O(n)$