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.PriorityQueue;

public class Solution {
    public long totalCost(int[] costs, int k, int candidates) {
        // Hint: Maintain two minheaps, one for the left end and one for the right end
        // This problem is intentionally made complex but actually we don't have to record the
        // indices
        int n = costs.length;
        PriorityQueue<Integer> leftMinHeap = new PriorityQueue<>();
        PriorityQueue<Integer> rightMinHeap = new PriorityQueue<>();
        long res = 0;
        if (2 * candidates >= n) {
            for (int i = 0; i <= n / 2; i++) {
                leftMinHeap.add(costs[i]);
            }
            for (int i = n / 2 + 1; i < n; i++) {
                rightMinHeap.add(costs[i]);
            }
            while (!leftMinHeap.isEmpty() && !rightMinHeap.isEmpty() && k > 0) {
                if (leftMinHeap.peek() <= rightMinHeap.peek()) {
                    res += leftMinHeap.poll();
                } else {
                    res += rightMinHeap.poll();
                }
                k -= 1;
            }
        } else {
            int left = candidates;
            int right = n - candidates - 1;
            for (int i = 0; i < candidates; i++) {
                leftMinHeap.add(costs[i]);
            }
            for (int i = n - candidates; i < n; i++) {
                rightMinHeap.add(costs[i]);
            }
            while (!leftMinHeap.isEmpty() && !rightMinHeap.isEmpty() && k > 0) {
                if (leftMinHeap.peek() <= rightMinHeap.peek()) {
                    res += leftMinHeap.poll();
                    if (left <= right) {
                        leftMinHeap.add(costs[left]);
                    }
                    left += 1;
                } else {
                    res += rightMinHeap.poll();
                    if (right >= left) {
                        rightMinHeap.add(costs[right]);
                    }
                    right -= 1;
                }
                k -= 1;
            }
        }
        if (k > 0 && leftMinHeap.isEmpty()) {
            while (k > 0) {
                res += rightMinHeap.poll();
                k -= 1;
            }
        }
        if (k > 0 && rightMinHeap.isEmpty()) {
            while (k > 0) {
                res += leftMinHeap.poll();
                k -= 1;
            }
        }
        return res;
    }
}