philosyang.com

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):

#Array #Greedy #Sorting #Heap #Python