mice-and-cheese.py (1149B)
1 import heapq 2 3 # two mice 4 # n types of cheese 5 # each type should be eaten by exactly one mouse 6 # reward1[i] if first mouse eats it 7 # reward2[i] if second mouse eats it 8 # k is non-negative, reward lists are all positive 9 # return max points if first mouse eats exactly k types of cheese 10 11 # idea: 12 # seems greedy 13 # like what's the delta at each position 14 # then minimize this somehow 15 16 # yes. Since first mouse eats only k cheese we compute the deltas and then select 17 # the ones that are largest wrt value of first eating - second eating 18 19 # this'll be O(n + klogk) for time 20 # We could do O(logk) for space 21 22 class Solution: 23 def miceAndCheese(self, reward1: List[int], reward2: List[int], k: int) -> int: 24 best_k_heap = [] 25 for i in range(len(reward1)): 26 delta = reward1[i] - reward2[i] 27 if len(best_k_heap) < k: 28 heapq.heappush(best_k_heap,delta) 29 else: 30 if k > 0 and best_k_heap[0] < delta: 31 heapq.heappop(best_k_heap) 32 heapq.heappush(best_k_heap,delta) 33 return sum(reward2) + sum(best_k_heap)