leetcode/algorithms/python3/2462-total-cost-to-hire-k-workers.py at main · remy727/leetcode · GitHub

Name: leetcode/algorithms/python3/2462-total-cost-to-hire-k-workers.py at main · remy727/leetcode · GitHub
Rating: 4.6 (2562 reviews)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
# 2462. Total Cost to Hire K Workers
# https://leetcode.com/problems/total-cost-to-hire-k-workers
# Medium

fromheapqimportheapify, heappush, heappop

classSolution:
deftotalCost(self, costs: list[int], k: int, candidates: int) ->int:
q=costs[:candidates]
qq=costs[max(candidates, len(costs)-candidates):]
heapify(q)
heapify(qq)
ans=0
i, ii=candidates, len(costs)-candidates-1
for_inrange(k):
ifnotqqorqandq[0] <=qq[0]:
ans+=heappop(q)
ifi<=ii:
heappush(q, costs[i])
i+=1
else:
ans+=heappop(qq)
ifi<=ii:
heappush(qq, costs[ii])
ii-=1
returnans

# ********************#
# TEST #
# ********************#

importunittest

classTestMinimumTime(unittest.TestCase):
deftest_minimumTime(self):
self.assertEqual(Solution.totalCost(self, [17,12,10,2,7,2,11,20,8], 3, 4), 11)
self.assertEqual(Solution.totalCost(self, [1,2,4,1], 3, 3), 4)

if__name__=='__main__':
unittest.main()