lintcode 1615 · 投资结果 【算法,堆】

题目

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;
    }
}

猜你喜欢

转载自blog.csdn.net/weixin_37991016/article/details/132783833