Problem statement

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

Solution

First of all, notice, that at each moment of time we try to solve the following problem: maximize $P_i$ given that $C_i \leqslant W$, where $W$ is current capital. Note, that if we take project with greater profit, then we increase $W$ and on the next steps our set of options will contain set of options if we choose smaller profit. Let us first sort Capitals $C_i$ and apply the same transform for Profits $P_i$. Then we will keep heap, where we keep all available pairs ($P_i, C_i$) (actually with minus sign, because in python heap is minimum heap). Now, we move through our Capitals one by one and add them to our heap, until it is possible, that is they are less than $W$. Finally, we extract (pop) the head of our heap: it will be the maximum gain on given step, which we add to $W$. Also it can happen, that we can not continue, because heap is empty, in this case we just return $W$ earlier.

Complexity

Time complexity is O(n log n): to sort our data, and then do at most n pushes to heap and at most k pops from heap. Space complexity is O(n).

Code

class Solution:
    def findMaximizedCapital(self, k, W, Profits, Capital):
        CP = sorted(zip(Capital,Profits))
        heap, itr, n = [], 0, len(CP)
        
        for _ in range(k):
            while itr < n and CP[itr][0] <= W:
                heappush(heap, (-CP[itr][1], -CP[itr][0]))
                itr += 1
            if not heap: return W
            prof, cap = heappop(heap)
            W -= prof
            
        return W