2462. Total Cost to Hire K Workers

Description

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:

  • You will run k sessions and hire exactly one worker in each session.
  • In each hiring session, choose the worker with the lowest cost from either the first candidates workers or the last candidates workers. Break the tie by the smallest index.
    • For example, if costs = [3,2,7,7,1,2] and candidates = 2, then in the first hiring session, we will choose the 4th worker because they have the lowest cost [3,2,7,7,1,2].
    • In the second hiring session, we will choose 1st worker because they have the same lowest cost as 4th worker but they have the smallest index [3,2,7,7,2]. Please note that the indexing may be changed in the process.
  • If there are fewer than candidates workers remaining, choose the worker with the lowest cost among them. Break the tie by the smallest index.
  • A worker can only be chosen once.

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.
- In the first hiring round we choose the worker from [17,12,10,2,7,2,11,20,8]. The lowest cost is 2, and we break the tie by the smallest index, which is 3. The total cost = 0 + 2 = 2.
- In the second hiring round we choose the worker from [17,12,10,7,2,11,20,8]. The lowest cost is 2 (index 4). The total cost = 2 + 2 = 4.
- In the third hiring round we choose the worker from [17,12,10,7,11,20,8]. The lowest cost is 7 (index 3). 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.
- In the first hiring round we choose the worker from [1,2,4,1]. The lowest cost is 1, and we break the tie by the smallest index, which is 0. The total cost = 0 + 1 = 1. Notice that workers with index 1 and 2 are common in the first and last 3 workers.
- In the second hiring round we choose the worker from [2,4,1]. The lowest cost is 1 (index 2). The total cost = 1 + 1 = 2.
- In the third hiring round there are less than three candidates. We choose the worker from the remaining workers [2,4]. The lowest cost is 2 (index 0). The total cost = 2 + 2 = 4.
The total hiring cost is 4.

 

Constraints:

  • 1 <= costs.length <= 105
  • 1 <= costs[i] <= 105
  • 1 <= k, candidates <= costs.length

Solutions

Solution 1: Priority Queue (Min Heap)

We maintain a priority queue (min heap) for the current candidate workers, and use variables $i$ and $j$ to mark the minimum index of the frontmost worker and the minimum index of the rearmost worker. Initially, $i = \text{candidates} - 1$ and $j = n - \text{candidates}$.

First, we put the costs of the first $candidates$ workers into the priority queue, then we put the costs of the last $candidates$ workers into the priority queue. Before putting them in, we need to check whether they are already in the priority queue according to $i$ or $j$. If they are, we don’t need to put them in again.

We loop $k$ times, each time taking out the worker with the smallest cost from the priority queue and accumulating the cost. If the index $x$ of the current worker is in the index range $[0,..i]$ of the frontmost workers, we move $i$ one step to the right, and then check whether we need to put the cost of the worker corresponding to $i$ into the priority queue; if the index is in the index range $[j,..n-1]$ of the rearmost workers, we move $j$ one step to the left, and then check whether we need to put the cost of the worker corresponding to $j$ into the priority queue.

After the traversal ends, we return the accumulated cost as the answer.

The time complexity is $O(n \times \log n)$, where $n$ is the length of the array $costs$.

Python Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

class Solution:
    def totalCost(self, costs: List[int], k: int, candidates: int) -> int:
        q = []
        n = len(costs)
        i, j = candidates - 1, n - candidates
        for h in range(candidates):
            q.append((costs[h], h))
        for h in range(n - candidates, n):
            if h > i:
                q.append((costs[h], h))
        heapify(q)
        ans = 0
        for _ in range(k):
            c, x = heappop(q)
            ans += c
            if x <= i:
                i += 1
                if i < j:
                    heappush(q, (costs[i], i))
            if x >= j:
                j -= 1
                if i < j:
                    heappush(q, (costs[j], j))
        return ans

Java Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

class Solution {
    public long totalCost(int[] costs, int k, int candidates) {
        PriorityQueue<int[]> q = new PriorityQueue<>((a, b) -> {
            if (a[0] == b[0]) {
                return a[1] - b[1];
            }
            return a[0] - b[0];
        });
        int n = costs.length;
        int i = candidates - 1, j = n - candidates;
        for (int h = 0; h < candidates; ++h) {
            q.offer(new int[] {costs[h], h});
        }
        for (int h = n - candidates; h < n; ++h) {
            if (h > i) {
                q.offer(new int[] {costs[h], h});
            }
        }
        long ans = 0;
        while (k-- > 0) {
            var e = q.poll();
            int c = e[0], x = e[1];
            ans += c;
            if (x <= i) {
                if (++i < j) {
                    q.offer(new int[] {costs[i], i});
                }
            }
            if (x >= j) {
                if (--j > i) {
                    q.offer(new int[] {costs[j], j});
                }
            }
        }
        return ans;
    }
}

C++ Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

using pii = pair<int, int>;

class Solution {
public:
    long long totalCost(vector<int>& costs, int k, int candidates) {
        priority_queue<pii, vector<pii>, greater<pii>> q;
        int n = costs.size();
        int i = candidates - 1, j = n - candidates;
        for (int h = 0; h < candidates; ++h) q.push({costs[h], h});
        for (int h = n - candidates; h < n; ++h)
            if (h > i) q.push({costs[h], h});
        long long ans = 0;
        while (k--) {
            auto [c, x] = q.top();
            q.pop();
            ans += c;
            if (x <= i) {
                if (++i < j) {
                    q.push({costs[i], i});
                }
            }
            if (x >= j) {
                if (--j > i) {
                    q.push({costs[j], j});
                }
            }
        }
        return ans;
    }
};

Go Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

func totalCost(costs []int, k int, candidates int) int64 {
	q := hp{}
	n := len(costs)
	i, j := candidates-1, n-candidates
	for h := 0; h < candidates; h++ {
		heap.Push(&q, pair{costs[h], h})
	}
	for h := n - candidates; h < n; h++ {
		if h > i {
			heap.Push(&q, pair{costs[h], h})
		}
	}
	ans := 0
	for k > 0 {
		p := heap.Pop(&q).(pair)
		c, x := p.c, p.x
		ans += c
		if x <= i {
			i++
			if i < j {
				heap.Push(&q, pair{costs[i], i})
			}
		}
		if x >= j {
			j--
			if i < j {
				heap.Push(&q, pair{costs[j], j})
			}
		}
		k--
	}
	return int64(ans)
}

type pair struct{ c, x int }
type hp []pair

func (h hp) Len() int           { return len(h) }
func (h hp) Less(i, j int) bool { return h[i].c < h[j].c || h[i].c == h[j].c && h[i].x < h[j].x }
func (h hp) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *hp) Push(v any)        { *h = append(*h, v.(pair)) }
func (h *hp) Pop() any          { a := *h; v := a[len(a)-1]; *h = a[:len(a)-1]; return v }