2611. Mice and Cheese | 1663
Am I greedy if I am good at greedy problems?
1class Solution:
2 def miceAndCheese(self, reward1: List[int], reward2: List[int], k: int) -> int:
3 # reward1 = [1, 1, 3, 4]
4 # reward2 = [4, 4, 1, 1]
5 # delta = [-3,-3,2, 3] # if eating reward1, how many points to get compared to eating reward2
6
7 n = len(reward1)
8 delta = []
9 heapq.heapify(delta)
10
11 for i in range(n):
12 # min heap so push inverted
13 heapq.heappush(delta, (reward2[i] - reward1[i], i))
14
15 reward1_index = set()
16
17 for j in range(k):
18 reward1_index.add(heapq.heappop(delta)[1])
19
20 result = 0
21 for x in range(n):
22 if x in reward1_index:
23 result += reward1[x]
24 else:
25 result += reward2[x]
26
27 return result
asymtotically optimal but there are improvables (thanks LLM):
- heapq.heapify(delta) on an empty list is unnecessary (no harm though).
- You don’t need to store indices or a set if you switch to the “base sum + diffs” pattern.
- Using heapq.nsmallest(k, delta) could also simplify the popping loop, but sorting is usually clearer.
#Array #Greedy #Sorting #Heap #Python