Total Cost to Hire K Workers

Sanjeev SharmaSanjeev Sharma
2 min read

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

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro