~/problems / Heaps / Heaps and priority queues

IPO

hard ~45 min

A small studio starts with a budget of start and may take on at most k projects, one after another. Project i can only be started if the current budget is at least needs[i] (the studio has to show it can cover the risk, but nothing is spent). Finishing it adds gains[i] to the budget. Each project can be done at most once.

Write max_capital(k, start, needs, gains) -> int that returns the largest budget the studio can end up with.

max_capital(2, 0, [0, 1, 1], [1, 2, 3])        # 4   (do project 0 -> 1, then project 2 -> 4)
max_capital(3, 0, [0, 1, 2], [1, 2, 3])        # 6   (0 -> 1 -> 3 -> 6)
max_capital(5, 1, [2, 3], [10, 10])            # 1   (nothing is affordable)
max_capital(1, 10, [5, 0, 10], [4, 7, 6])      # 17

Constraints:

  • n = len(needs) == len(gains), 1 <= n <= 10^5 and 1 <= k <= 10^5.
  • 0 <= start, needs[i], gains[i] <= 10^9. The final budget can exceed 32-bit range.
  • Rescanning every project before each pick costs O(n·k), far too slow when both are 10^5. Aim for O((n + k) log n).
Show hint

Your budget never goes down, so once a project becomes affordable it stays affordable. Walk through the projects in order of what they need, move the newly affordable ones into a pool, and always take the pool's most rewarding one.

Topic: Heaps and priority queues. heapq: top-k, k-way merge, two heaps for a running median.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc