LeetCode 502: IPO

link

Greedy

Among the projects we have enough capital to do, we pick the most profitable one.

  1. Keeping capital sorted in increasing order helps find projects we have enough capital to do. So, we use a min-heap with capital as key.
  2. Among the projects from (1) we want to pick the one with maximum profit. So, we use a max-heap with profit as key.

Since total capital never decreases, once a project has been moved to the profit heap, unless we decide to do it, it remains there. Thus, a project is either in capital heap or profit heap. Each project is pushed to and popped from the capital heap at most once and same with profit heap.

Time: \mathcal{O}(n \cdot \lg{n}), space: \mathcal{O}(n).

from heapq import heapify, heappush as push, heappop as pop

class Solution:
    def findMaximizedCapital(
        self, k: int, w: int, profits: List[int], capital: List[int]
    ) -> int:
        capq = [(c, i) for i, c in enumerate(capital)]
        heapify(capq)
        profq = []
        for _ in range(k):
            while capq:
                if w < capq[0][0]:
                    break
                _, i = pop(capq)
                push(profq, -profits[i])
            if not profq:
                break
            w += -pop(profq)

        return w

Leave a comment