IPO LeetCode solution in Python

1
from heapq import heappush, heappop, nlargest
2

3

4
class Solution:
5
    def findMaximizedCapital(
6
        self, k: int, w: int, profits: List[int], capital: List[int]
7
    ) -> int:
8
        if w >= max(capital):
9
            return w + sum(nlargest(k, profits))
10

11
        projects = [[capital[i], profits[i]] for i in range(len(profits))]
12
        projects.sort(key=lambda x: x[0])
13

14
        heap = []
15

16
        for i in range(k):
17
            while projects and projects[0][0] <= w:
18
                heappush(heap, -1 * projects.pop(0)[1])
19

20
            if not heap:
21
                break
22
            p = -heappop(heap)
23
            w += p
24
        return w