heap
greedy
]
Leetcode 0857 Minimum Cost to Hire K Workers
Problem statement
https://leetcode.com/problems/minimum-cost-to-hire-k-workers/
Solution
Denote $q_0, \dots, q_{n-1}$ qualities and by $w_0, \dots w_{n-1}$ wages. Then our goal is to find such indexes $i_1,\dots i_k$, for which \(\alpha (q_{i_1} + \dots + q_{i_k})\) is minimized, where $\alpha\cdot q_{i_1} \geqslant w_{i_1}, \dots \alpha\cdot q_{i_k} \geqslant w_{i_k}$. Let us call $\alpha$ level of salary: so here we need to chose $\alpha = \max \limits_{j = 1\dots k} w_{i_j}/q_{i_j}$: so we can see that there will be at most $n$ different levels of salary we can choose.
Now, let us sort our data by levels of salary: $s_0 \leqslant s_1 \leqslant \dots \leqslant s_{n-1}$.
If we choose level $\alpha_0$, then we can hire only $1$ person: with index $0$, if we choose level $\alpha_1$, then we can hire $2$ persons: with indexes $0$ and $1$ and so on. We can start with index $k-1$, where we can hire $k$ persons. Now, imagine we reached level $s_l$, $l\geqslant k$ of salary: then we can hire $l+1$ first persons, how we choose, which ones? We need to take persons with smallest qualities (see equation, $\alpha$ is fixed). So, we will keep max heap with k
smallest qualities, and on each step, when new person arrives, we update it: if new element is bigger than root of heap, we do nothing, if it is smaller, we push it to heap and remove root in O(log k)
. Also we recalculate sum of element in heap in O(1)
. Finally, we update our ans
.
Complexity
Time complexity is O(n log k)
, space complexity is O(n)
.
Code
class Solution:
def mincostToHireWorkers(self, quality, wage, K):
nums = sorted((w/q, q) for q, w in zip(quality, wage))
n, heap = len(nums), []
for i in range(K): heappush(heap, -nums[i][1])
sum_q, level = -sum(heap), nums[K-1][0]
ans = sum_q * level
for i in range(K, n):
if nums[i][1] <= - heap[0]:
heappush(heap, -nums[i][1])
t = heappop(heap)
sum_q = sum_q + t + nums[i][1]
ans = min(ans, nums[i][0] * sum_q)
return ans