這題算是典型的 greedy 題,我們每次都想要投資目前 capital 能買到的標的中,profit 最大的那一個,但因為我們可以投資不只一個 project,所以每次投資完,我們都希望能夠最快計算出下一個投資標的。要做到這件事,用兩個 heap 再適合不過了。
1st - max heap:儲存目前買得起的 project,儲存 ,把 profit 最高的排在最上面。 2nd - min heap:儲存目前買不起的 project,儲存 ,把 capital 最小的排在最上面。
每次投資前,我們只要從 1st heap 取出第一個,更新完 capital,把 2nd heap 新買得起的 project 放進 1st heap,就可以開始下一輪投資。
using namespace std;
#include <iostream>
#include <queue>
#include <vector>
class MaximizeCapital {
public:
struct comp1 {
bool operator() (const pair<int, int>& p1, const pair<int, int>& p2) {
return p1.second < p2.second;
}
};
struct comp2 {
bool operator() (const pair<int, int>& p1, const pair<int, int>& p2) {
return p1.first > p2.first;
}
};
static int findMaximumCapital(const vector<int> &capital, const vector<int> &profits,
int numberOfProjects, int initialCapital) {
// Create two heaps
priority_queue<pair<int, int>, vector<pair<int, int>>, comp1> maxHeap;
priority_queue<pair<int, int>, vector<pair<int, int>>, comp2> minHeap;
for(int i = 0; i < capital.size(); ++i) {
if(capital[i] <= initialCapital) {
maxHeap.push(make_pair(capital[i], profits[i]));
}
else {
minHeap.push(make_pair(capital[i], profits[i]));
}
}
// Start investing
int profit = initialCapital;
while(numberOfProjects--) {
pair<int, int> cur = maxHeap.top();
maxHeap.pop();
profit += cur.second;
initialCapital += cur.second;
while(!minHeap.empty() && minHeap.top().first <= initialCapital) {
maxHeap.push(minHeap.top());
minHeap.pop();
}
}
return profit;
}
};
int main(int argc, char *argv[]) {
int result =
MaximizeCapital::findMaximumCapital(vector<int>{0, 1, 2}, vector<int>{1, 2, 3}, 2, 1);
cout << "Maximum capital: " << result << endl;
result =
MaximizeCapital::findMaximumCapital(vector<int>{0, 1, 2, 3}, vector<int>{1, 2, 3, 5}, 3, 0);
cout << "Maximum capital: " << result << endl;
}