LC-6231.雇佣K位工人的总代价

题目描述

leetcode 中等题

给你一个下标从 0 开始的整数数组 costs ，其中 costs[i] 是雇佣第 i 位工人的代价。

同时给你两个整数 k 和 candidates 。我们想根据以下规则恰好雇佣 k 位工人：

总共进行 k 轮雇佣，且每一轮恰好雇佣一位工人。
在每一轮雇佣中，从最前面 candidates 和最后面 candidates 人中选出代价最小的一位工人，如果有多位代价相同且最小的工人，选择下标更小的一位工人。
比方说，costs = [3,2,7,7,1,2] 且 candidates = 2 ，第一轮雇佣中，我们选择第 4 位工人，因为他的代价最小 [3,2,7,7,1,2] 。
第二轮雇佣，我们选择第 1 位工人，因为他们的代价与第 4 位工人一样都是最小代价，而且下标更小，[3,2,7,7,2] 。注意每一轮雇佣后，剩余工人的下标可能会发生变化。
如果剩余员工数目不足 candidates 人，那么下一轮雇佣他们中代价最小的一人，如果有多位代价相同且最小的工人，选择下标更小的一位工人。
一位工人只能被选择一次。
返回雇佣恰好 k 位工人的总代价。

示例1：
输入：costs = [17,12,10,2,7,2,11,20,8], k = 3, candidates = 4
输出：11
解释：我们总共雇佣 3 位工人。总代价一开始为 0 。

• 第一轮雇佣，我们从 [17,12,10,2,,2,11,20,8] 中选择。最小代价是 2 ，有两位工人，我们选择下标更小的一位工人，即第 3 位工人。总代价是 0 + 2 = 2 。
• 第二轮雇佣，我们从 [17,12,10,7,2,11,20,8] 中选择。最小代价是 2 ，下标为 4 ，总代价是 2 + 2 = 4 。
• 第三轮雇佣，我们从 [17,12,10,7,11,20,8] 中选择，最小代价是 7 ，下标为 3 ，总代价是 4 + 7 = 11 。注意下标为 3 的工人同时在最前面和最后面 4 位工人中。
  总雇佣代价是 11 。

提示：

1 <= costs.length <= 10^5
1 <= costs[i] <= 10^5
1 <= k, candidates <= costs.length

最小堆

通过两个最小堆模拟 最前面 candidates 和最后面 candidates 个人 即可，具体来说每次取出两个堆的最小值，接着继续尝试维护堆直到 k 轮选人结束。

class Solution {
    public long totalCost(int[] costs, int k, int candidates) {
        int n = costs.length;
        Comparator<int[]> cmp = (a, b) -> {
            if(a[0] == b[0]){
                return a[1] - b[1];
            }
            return a[0] - b[0];
        };
        // 实际上不需要比较下标，可以改成 PriorityQueue<Integer>
        PriorityQueue<int[]> front = new PriorityQueue<>(cmp); 
        PriorityQueue<int[]> tail = new PriorityQueue<>(cmp);
        int left = 0, right = n - 1;
        for(int i = 0; i < candidates && left <= right; i++){ // candidates 个人
            front.offer(new int[]{costs[left], left++});
            if(left > right){
                break;
            }
            tail.offer(new int[]{costs[right], right--});
        }
        long ans = 0;
        while(k-- > 0){
            int[] f = front.peek();
            int[] t = tail.peek();
            if(t == null || (f != null && f[0] <= t[0])){
                ans += front.poll()[0];
                if(left <= right){
                    front.offer(new int[]{costs[left], left++});
                }
            }else{
                ans += tail.poll()[0];
                if(left <= right){
                    tail.offer(new int[]{costs[right], right--});
                }
            }
        }
        return ans;
    }
}