給定一個數(shù)列costs,一共k次操作,每次從前candidates和后candidates個數(shù)里面選較小的一個,問k次取到的k個數(shù)只和
維護兩個優(yōu)先隊列,一個存前candidates,一個存后candidates,注意每次pop之后要push后一個/前一個數(shù)
1 #2462
2 #Runtime: 1356 ms (Beats 44.93%)
3 #Memory: 21.8 MB (Beats 68.12%)
4
5 class Solution(object):
6 def totalCost(self, costs, k, candidates):
7 """
8 :type costs: List[int]
9 :type k: int
10 :type candidates: int
11 :rtype: int
12 """
13 left = costs[:candidates]
14 right = costs[max(candidates, len(costs) - candidates):]
15 heapify(left)
16 heapify(right)
17 ans = 0
18 i = candidates
19 j = len(costs) - candidates - 1
20 for _ in range(k):
21 if (not right) or (left and left[0] <= right[0]):
22 ans += heappop(left)
23 if i <= j:
24 heappush(left, costs[i])
25 i += 1
26 else:
27 ans += heappop(right)
28 if i <= j:
29 heappush(right, costs[j])
30 j -= 1
31 return ans