LeetCode in Kotlin

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

@Suppress("NAME_SHADOWING")
class Solution {
    fun totalCost(costs: IntArray, k: Int, candidates: Int): Long {
        // 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
        var k = k
        val n = costs.size
        val leftMinHeap = PriorityQueue<Int>()
        val rightMinHeap = PriorityQueue<Int>()
        var res: Long = 0
        if (2 * candidates >= n) {
            for (i in 0..n / 2) {
                leftMinHeap.add(costs[i])
            }
            for (i in n / 2 + 1 until n) {
                rightMinHeap.add(costs[i])
            }
            while (leftMinHeap.isNotEmpty() && rightMinHeap.isNotEmpty() && k > 0) {
                res += if (leftMinHeap.peek() <= rightMinHeap.peek()) {
                    leftMinHeap.poll().toLong()
                } else {
                    rightMinHeap.poll().toLong()
                }
                k -= 1
            }
        } else {
            var left = candidates
            var right = n - candidates - 1
            for (i in 0 until candidates) {
                leftMinHeap.add(costs[i])
            }
            for (i in n - candidates until n) {
                rightMinHeap.add(costs[i])
            }
            while (leftMinHeap.isNotEmpty() && rightMinHeap.isNotEmpty() && k > 0) {
                if (leftMinHeap.peek() <= rightMinHeap.peek()) {
                    res += leftMinHeap.poll().toLong()
                    if (left <= right) {
                        leftMinHeap.add(costs[left])
                    }
                    left += 1
                } else {
                    res += rightMinHeap.poll().toLong()
                    if (right >= left) {
                        rightMinHeap.add(costs[right])
                    }
                    right -= 1
                }
                k -= 1
            }
        }
        if (k > 0 && leftMinHeap.isEmpty()) {
            while (k > 0) {
                res += rightMinHeap.poll().toLong()
                k -= 1
            }
        }
        if (k > 0 && rightMinHeap.isEmpty()) {
            while (k > 0) {
                res += leftMinHeap.poll().toLong()
                k -= 1
            }
        }
        return res
    }
}