857. Minimum Cost to Hire K Workers

[zeikar]

Problem link

https://leetcode.com/problems/minimum-cost-to-hire-k-workers/

Problem Summary

노동자의 품질과 최소 임금이 주어질 때 K 명의 노동자를 뽑아야 한다. K명의 노동자가 받는 임금과 품질의 비율은 모두 같아야 한다.
이때 K 명을 뽑을 수 있는 최소 임금을 구하는 문제.

Solution

처음부터 드는 생각은 가성비가 좋은 노동자를 뽑아보는 것이다. 즉, wage / quality 를 한 값이 작은 순서로 K명을 뽑아주면 가성비가 좋은 K 명을 뽑을 수 있다.
단, 전체 임금의 합의 최소를 구하는 문제이므로 가성비가 좋지 않아도 임금이 더 적게 든다면 해당 노동자를 뽑는 게 낫다.

예제 1번이 그 예시인데, 두 번째 노동자가 quality 20으로 가성비가 좋지만 그만큼 많은 임금이 필요하므로 선택하지 않는 것이 더 최소가 된다.

구현은 가성비로 정렬을 한 후 돌면서 K명을 뽑았다면 result에 업데이트 해주고 K명이 넘어가면 힙에서 q 값이 큰 순서로 뽑아주면서 (max heap) 갱신하면 된다.

참고: https://leetcode.com/problems/minimum-cost-to-hire-k-workers/discuss/141768/Detailed-explanation-O(NlogN)

Source Code

import heapq
from typing import List


class Solution:
    def mincostToHireWorkers(self, quality: List[int], wage: List[int], k: int) -> float:
        workers = sorted([float(w) / q, q] for w, q in zip(wage, quality))
        result = float('inf')
        qsum = 0
        heap = []

        for r, q in workers:
            heapq.heappush(heap, -q)
            qsum += q

            if len(heap) > k:
                qsum += heapq.heappop(heap)

            if len(heap) == k:
                result = min(result, qsum * r)
        return result