Problem of The Day: Total Cost to Hire K Workers

Problem Statement

leetcode problem link

My Solution [TLE]

class Solution:
    def totalCost(self, costs: List[int], k: int, candidates: int) -> int:
        res = 0
        cand_count = 0
        while cand_count < k:
            min_heap = []
            N = len(costs)
            for i in range(min(candidates, N)):
                heapq.heappush(min_heap, (costs[i], i))

            cand_num = 0
            for j in range(len(costs) - 1, -1, -1):
                heapq.heappush(min_heap, (costs[j], j))
                cand_num += 1
                if cand_num == candidates:
                    break

            cost, i = heapq.heappop(min_heap)
            costs.pop(i)
            res += cost
            cand_count += 1

        return res

Editorial

Approach 1: 2 Priority Queues

class Solution:
    def totalCost(self, costs: List[int], k: int, candidates: int) -> int:
        # head_workers stores the first k workers.
        # tail_workers stores at most last k workers without any workers from the first k workers.
        head_workers = costs[:candidates]
        tail_workers = costs[max(candidates, len(costs) - candidates):]
        heapify(head_workers)
        heapify(tail_workers)

        answer = 0
        next_head, next_tail = candidates, len(costs) - 1 - candidates

        for _ in range(k):
            if not tail_workers or head_workers and head_workers[0] <= tail_workers[0]:
                answer += heappop(head_workers)

                # Only refill the queue if there are workers outside the two queues.
                if next_head <= next_tail:
                    heappush(head_workers, costs[next_head])
                    next_head += 1
            else:
                answer += heappop(tail_workers)

                # Only refill the queue if there are workers outside the two queues.
                if next_head <= next_tail:
                    heappush(tail_workers, costs[next_tail])
                    next_tail -= 1

        return answer

Approach 2: 1 Priority Queue

class Solution:
    def totalCost(self, costs: List[int], k: int, candidates: int) -> int:
        # Add the first k workers with section id of 0 and
        # the last k workers with section id of 1 (without duplication) to pq.
        pq = []
        for i in range(candidates):
            pq.append((costs[i], 0))
        for i in range(max(candidates, len(costs) - candidates), len(costs)):
            pq.append((costs[i], 1))

        heapify(pq)

        answer = 0
        next_head, next_tail = candidates, len(costs) - 1 - candidates

        # Only refill pq if there are workers outside.
        for _ in range(k):
            cur_cost, cur_section_id = heappop(pq)
            answer += cur_cost
            if next_head <= next_tail:
                if cur_section_id == 0:
                    heappush(pq, (costs[next_head], 0))
                    next_head += 1
                else:
                    heappush(pq, (costs[next_tail], 1))
                    next_tail -= 1

        return answer