Total Cost to Hire K Workers
Advertisement
Problem
Hire k workers. Each round pick cheapest among first p and last p candidates (tie: pick smallest index).
Key insight: Two min-heaps: one for left p candidates, one for right p candidates. Expand as we hire.
Solutions
# Python
import heapq
def totalCost(costs, k, candidates):
n = len(costs)
left = list(range(candidates))
right = list(range(max(candidates, n - candidates), n))
left_heap = [(costs[i], i) for i in range(min(candidates, n))]
right_heap = [(costs[i], i) for i in range(max(n - candidates, candidates), n)]
heapq.heapify(left_heap)
heapq.heapify(right_heap)
l, r = candidates, n - candidates - 1
total = 0
for _ in range(k):
lv = left_heap[0] if left_heap else (float('inf'), -1)
rv = right_heap[0] if right_heap else (float('inf'), -1)
if lv[0] <= rv[0]:
c, i = heapq.heappop(left_heap)
total += c
if l <= r:
heapq.heappush(left_heap, (costs[l], l))
l += 1
else:
c, i = heapq.heappop(right_heap)
total += c
if l <= r:
heapq.heappush(right_heap, (costs[r], r))
r -= 1
return total
// C++
long long totalCost(vector<int>& costs, int k, int candidates) {
int n=costs.size();
priority_queue<pair<int,int>,vector<pair<int,int>>,greater<>> lH,rH;
int l=0,r=n-1;
for(int i=0;i<candidates&&l<=r;i++,l++) lH.push({costs[l],l});
for(int i=0;i<candidates&&l<=r;i++,r--) rH.push({costs[r],r});
long long total=0;
for(int i=0;i<k;i++){
bool useLeft=!lH.empty()&&(rH.empty()||lH.top()<=rH.top());
auto&h=useLeft?lH:rH;
auto[c,idx]=h.top();h.pop(); total+=c;
if(l<=r){if(useLeft){lH.push({costs[l],l});l++;}else{rH.push({costs[r],r});r--;}}
}
return total;
}
// Java
public long totalCost(int[] costs, int k, int candidates) {
int n=costs.length;
PriorityQueue<int[]>lH=new PriorityQueue<>((a,b)->a[0]!=b[0]?a[0]-b[0]:a[1]-b[1]);
PriorityQueue<int[]>rH=new PriorityQueue<>((a,b)->a[0]!=b[0]?a[0]-b[0]:a[1]-b[1]);
int l=0,r=n-1;
for(int i=0;i<candidates&&l<=r;i++,l++) lH.offer(new int[]{costs[l],l});
for(int i=0;i<candidates&&l<=r;i++,r--) rH.offer(new int[]{costs[r],r});
long total=0;
for(int i=0;i<k;i++){
boolean useL=!lH.isEmpty()&&(rH.isEmpty()||lH.peek()[0]<rH.peek()[0]||(lH.peek()[0]==rH.peek()[0]&&lH.peek()[1]<=rH.peek()[1]));
int[]p=(useL?lH:rH).poll(); total+=p[0];
if(l<=r){if(useL){lH.offer(new int[]{costs[l],l});l++;}else{rH.offer(new int[]{costs[r],r});r--;}}
}
return total;
}
// JavaScript
function totalCost(costs, k, candidates) {
const n = costs.length;
const lH = [], rH = [];
let l = 0, r = n - 1;
const push_sorted = (h, item) => { let p = h.findIndex(x => x[0] > item[0] || (x[0]===item[0]&&x[1]>item[1])); if(p===-1)h.push(item);else h.splice(p,0,item); };
for(let i=0;i<candidates&&l<=r;i++,l++) push_sorted(lH,[costs[l],l]);
for(let i=0;i<candidates&&l<=r;i++,r--) push_sorted(rH,[costs[r],r]);
let total = 0;
for(let i=0;i<k;i++){
const lv=lH[0]||[Infinity,-1], rv=rH[0]||[Infinity,-1];
if(lv[0]<rv[0]||(lv[0]===rv[0]&&lv[1]<=rv[1])){total+=lH.shift()[0];if(l<=r){push_sorted(lH,[costs[l],l]);l++;}}
else{total+=rH.shift()[0];if(l<=r){push_sorted(rH,[costs[r],r]);r--;}}
}
return total;
}
Complexity
- Time: O((k + candidates) log candidates)
- Space: O(candidates)
Advertisement
← Previous
Seat Reservation Manager