Skip to content

2611. Mice and Cheese 👍

  • Time: C++: $O(n)$, Java/Python: $O(n\log k)$
  • Space: C++: $O(1)$, Java/Python: $O(k)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class Solution {
 public:
  int miceAndCheese(vector<int>& reward1, vector<int>& reward2, int k) {
    // diffs[i] := reward1[i] - reward2[i].
    vector<int> diffs;

    for (int i = 0; i < reward1.size(); ++i)
      diffs.push_back(reward1[i] - reward2[i]);

    nth_element(diffs.begin(), diffs.begin() + k, diffs.end(), greater<>());
    return accumulate(reward2.begin(), reward2.end(), 0) +
           accumulate(diffs.begin(), diffs.begin() + k, 0);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution {
  public int miceAndCheese(int[] reward1, int[] reward2, int k) {
    Queue<Integer> minHeap = new PriorityQueue<>();

    for (int i = 0; i < reward1.length; ++i) {
      minHeap.offer(reward1[i] - reward2[i]);
      if (minHeap.size() == k + 1)
        minHeap.poll();
    }

    return Arrays.stream(reward2).sum() + minHeap.stream().mapToInt(Integer::intValue).sum();
  }
}

1
2
3
class Solution:
  def miceAndCheese(self, reward1: List[int], reward2: List[int], k: int) -> int:
    return sum(reward2) + sum(heapq.nlargest(k, (a - b for a, b in zip(reward1, reward2))))