Skip to content

502. IPO 👍

  • Time: $O(n\log n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
struct T {
  int pro;
  int cap;
};

class Solution {
 public:
  int findMaximizedCapital(int k, int w, vector<int>& profits,
                           vector<int>& capital) {
    auto compareC = [](const T& a, const T& b) { return a.cap > b.cap; };
    auto compareP = [](const T& a, const T& b) { return a.pro < b.pro; };
    priority_queue<T, vector<T>, decltype(compareC)> minHeap(compareC);
    priority_queue<T, vector<T>, decltype(compareP)> maxHeap(compareP);

    for (int i = 0; i < capital.size(); ++i)
      minHeap.emplace(profits[i], capital[i]);

    while (k-- > 0) {
      while (!minHeap.empty() && minHeap.top().cap <= w)
        maxHeap.push(minHeap.top()), minHeap.pop();
      if (maxHeap.empty())
        break;
      w += maxHeap.top().pro, maxHeap.pop();
    }

    return w;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
  public int findMaximizedCapital(int k, int w, int[] profits, int[] capital) {
    record T(int pro, int cap) {}
    Queue<T> minHeap = new PriorityQueue<>((a, b) -> Integer.compare(a.cap, b.cap));
    Queue<T> maxHeap = new PriorityQueue<>((a, b) -> Integer.compare(b.pro, a.pro));

    for (int i = 0; i < capital.length; ++i)
      minHeap.offer(new T(profits[i], capital[i]));

    while (k-- > 0) {
      while (!minHeap.isEmpty() && minHeap.peek().cap <= w)
        maxHeap.offer(minHeap.poll());
      if (maxHeap.isEmpty())
        break;
      w += maxHeap.poll().pro;
    }

    return w;
  }
}