LeetCode-in-Java

2462. Total Cost to Hire K Workers

Medium

You are given a 0-indexed integer array costs where costs[i] is the cost of hiring the ith worker.

You are also given two integers k and candidates. We want to hire exactly k workers according to the following rules:

Return the total cost to hire exactly k workers.

Example 1:

Input: costs = [17,12,10,2,7,2,11,20,8], k = 3, candidates = 4

Output: 11

Explanation: We hire 3 workers in total. The total cost is initially 0.

The total cost = 4 + 7 = 11. Notice that the worker with index 3 was common in the first and last four workers.

The total hiring cost is 11.

Example 2:

Input: costs = [1,2,4,1], k = 3, candidates = 3

Output: 4

Explanation: We hire 3 workers in total. The total cost is initially 0.

The total cost = 2 + 2 = 4. The total hiring cost is 4.

Constraints:

Solution

import java.util.Arrays;
import java.util.PriorityQueue;

public class Solution {
    public long totalCost(int[] costs, int k, int candidates) {
        int n = costs.length;
        long res = 0;
        if (candidates * 2 + k > n) {
            Arrays.sort(costs);
            for (int i = 0; i < k; i++) {
                res += costs[i];
            }
            return res;
        }
        PriorityQueue<Integer> pre = new PriorityQueue<>();
        PriorityQueue<Integer> suf = new PriorityQueue<>();
        for (int i = 0; i < candidates; i++) {
            pre.offer(costs[i]);
            suf.offer(costs[n - 1 - i]);
        }
        int i = candidates;
        int j = n - 1 - candidates;
        while (k-- > 0) {
            if (pre.peek() <= suf.peek()) {
                res += pre.poll();
                pre.offer(costs[i++]);
            } else {
                res += suf.poll();
                suf.offer(costs[j--]);
            }
        }
        return res;
    }
}