题目
https://www.lintcode.com/problem/1615
给定一个列表funds表示投资人每次的投资额。现在有三个公司A, B, C,它们的初始资金分别为a,b,c。投资人每次投资时会对当前资金最少的公司进行投资(当有多个公司资金相同时,投资人会对编号最小的公司进行投资)。返回A, B, C三家公司最后的资金。
1<= funds的长度<= 500000
1<= funds[i],a,b<=100
样例
样例1
输入: funds=[1,2,1,3,1,1], a=1, b=2, c=1
输出: [4,5,4]
解释:
第一轮投资中,A,C投资量一样,所以我们选A,此时a=2,b=2,c=1
第二轮投资中,C投资量最少,所以我们选C,此时a=2,b=2,c=3
第三轮投资中,A,B投资量一样,所以我们选A,此时a=3,b=2,c=3
第四轮投资中,B投资量最少,所以我们选B,此时a=3,b=5,c=3
第五轮投资中,A,C投资量一样,所以我们选A,此时a=4,b=5,c=3
第六轮投资中,C投资量最少,所以我们选C,此时a=4,b=5,c=4
样例2
输入: funds=[2,1,1,1], a=1, b=2, c=2
输出: [4,3,3]
解释:
第一轮投资中,A投资量最少,所以我们选A,此时a=3,b=2,c=2
第二轮投资中,B,C投资量一样,所以我们选B,此时a=3,b=3,c=2
第三轮投资中,C投资量最少,所以我们选C,此时a=3,b=3,c=3
第四轮投资中,A,B,C投资量一样,所以我们选A,此时a=4,b=3,c=3
思路
堆,又叫优先级队列类解决
参考代码
public class Solution {
/**
* @param funds: The investment each time
* @param a: The initial funds of A
* @param b: The initial funds of B
* @param c: The initial funds of C
* @return: The final funds
*/
public int[] getAns(int[] funds, int a, int b, int c) {
PriorityQueue<int[]> q = new PriorityQueue<>(new Comparator<int[]>() {
@Override
public int compare(int[] a, int[] b) {
return a[0] == b[0] ? a[1] - b[1] : a[0] - b[0];
}
});
q.add(new int[]{
a, 1}); //加入a
q.add(new int[]{
b, 2}); //加入b
q.add(new int[]{
c, 3}); // 加入c
int n = funds.length;
for (int i = 0; i < n; i++) {
int[] arr = q.poll();
int sum = arr[0], index = arr[1];
sum += funds[i];
q.add(new int[]{
sum, index});
}
int[] ans = new int[3];
while (!q.isEmpty()) {
int[] arr = q.poll();
int sum = arr[0];
int index = arr[1];
ans[index - 1] = sum;
}
return ans;
}
}