2611. Mice and Cheese

MediumGreedyArraySortingHeap (Priority Queue)

Problem Description

You have two mice and n different types of cheese. Each piece of cheese must be eaten by exactly one mouse.

For each cheese at index i (0-indexed):

If the first mouse eats it, you get reward1[i] points
If the second mouse eats it, you get reward2[i] points

Given:

reward1: an array of positive integers representing points if the first mouse eats each cheese
reward2: an array of positive integers representing points if the second mouse eats each cheese
k: a non-negative integer representing exactly how many pieces of cheese the first mouse must eat

Your task is to find the maximum total points possible when the first mouse eats exactly k pieces of cheese and the second mouse eats the remaining n - k pieces.

For example, if you have 4 types of cheese with reward1 = [1, 4, 4, 6], reward2 = [6, 5, 3, 6], and k = 1, the optimal strategy would be for the first mouse to eat cheese at index 2 (gaining 4 points) and the second mouse to eat the rest (gaining 6 + 5 + 6 = 17 points), for a total of 21 points.

The solution uses a greedy approach by calculating the difference reward1[i] - reward2[i] for each cheese. This represents the net gain (or loss) of giving cheese i to the first mouse instead of the second mouse. By sorting these differences in descending order and selecting the top k cheeses for the first mouse, we maximize the total score.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

Let's think about this problem step by step. We need to distribute n pieces of cheese between two mice, with the first mouse getting exactly k pieces.

One way to approach this is to start by assuming we give all cheese to the second mouse. This gives us a baseline score of sum(reward2). Now, we need to select k pieces to reassign from the second mouse to the first mouse.

When we reassign cheese i from the second mouse to the first mouse:

We lose reward2[i] points (since the second mouse no longer eats it)
We gain reward1[i] points (since the first mouse now eats it)
The net change is reward1[i] - reward2[i]

To maximize our total score, we want to reassign the k pieces that give us the largest positive net change. Even if some changes are negative, we still need to pick the k least negative ones since the first mouse must eat exactly k pieces.

This leads us to the key insight: sort all cheeses by their net change value reward1[i] - reward2[i] in descending order. The first k cheeses in this sorted order should go to the first mouse, as they provide the maximum benefit (or minimum loss) when reassigned.

For example, if the differences are [3, -2, 5, -1] and k = 2, we'd sort to get [5, 3, -1, -2] and give the cheeses with differences 5 and 3 to the first mouse, maximizing our total gain.

This greedy strategy works because each cheese assignment decision is independent - the benefit of assigning cheese i to a particular mouse doesn't depend on which other cheeses are assigned to which mice.

Learn more about Greedy, Sorting and Heap (Priority Queue) patterns.

Solution Approach

Based on our greedy strategy, here's how we implement the solution:

Calculate the differences: We need to sort the cheese indices based on the value reward1[i] - reward2[i]. This difference represents the net gain when cheese i is given to the first mouse instead of the second mouse.
Sort indices by difference: Instead of sorting the arrays directly, we sort the indices. We create a list of indices from 0 to n-1 and sort them based on the difference values in descending order:
```
idx = sorted(range(n), key=lambda i: reward1[i] - reward2[i], reverse=True)
```
This gives us the order in which we should prioritize giving cheese to the first mouse.
Distribute the cheese:
- The first k indices in our sorted list represent the cheese that should go to the first mouse (these have the highest difference values)
- The remaining n - k indices represent the cheese that should go to the second mouse

Calculate the total score:

return sum(reward1[i] for i in idx[:k]) + sum(reward2[i] for i in idx[k:])

Sum up reward1[i] for the first k indices (cheese eaten by the first mouse)
Sum up reward2[i] for the remaining indices (cheese eaten by the second mouse)

Time Complexity: O(n log n) due to sorting the indices
Space Complexity: O(n) for storing the sorted indices

The beauty of this approach is that by sorting based on the difference reward1[i] - reward2[i], we ensure that we're making the globally optimal choice. The cheese with the highest difference values are the ones where the first mouse has the greatest advantage (or smallest disadvantage) over the second mouse, so these should be prioritized for the first mouse.

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example to illustrate the solution approach.

Given:

reward1 = [1, 4, 4, 6, 4]
reward2 = [6, 5, 3, 6, 1]
k = 2 (first mouse must eat exactly 2 pieces)

Step 1: Calculate differences

For each cheese, calculate reward1[i] - reward2[i]:

Cheese 0: 1 - 6 = -5 (giving to mouse 1 loses 5 points vs mouse 2)
Cheese 1: 4 - 5 = -1 (giving to mouse 1 loses 1 point vs mouse 2)
Cheese 2: 4 - 3 = 1 (giving to mouse 1 gains 1 point vs mouse 2)
Cheese 3: 6 - 6 = 0 (no difference which mouse eats it)
Cheese 4: 4 - 1 = 3 (giving to mouse 1 gains 3 points vs mouse 2)

Step 2: Sort indices by difference (descending)

Differences with indices: [(0, -5), (1, -1), (2, 1), (3, 0), (4, 3)]

After sorting by difference value (highest first):

Index 4 with difference 3
Index 2 with difference 1
Index 3 with difference 0
Index 1 with difference -1
Index 0 with difference -5

Sorted indices: [4, 2, 3, 1, 0]

Step 3: Distribute cheese

Since k = 2, the first mouse gets the top 2 cheeses from our sorted list:

First mouse eats: cheese 4 and cheese 2
Second mouse eats: cheese 3, cheese 1, and cheese 0

Step 4: Calculate total score

First mouse score: reward1[4] + reward1[2] = 4 + 4 = 8 Second mouse score: reward2[3] + reward2[1] + reward2[0] = 6 + 5 + 6 = 17 Total score: 8 + 17 = 25

Verification: By choosing the cheeses with the highest differences (3 and 1), we maximized our gain. Even though cheese 2's difference is only 1, it's better than giving the first mouse cheese 3 (difference 0) or worse, cheese 1 or 0 (negative differences).

Solution Implementation

1class Solution:
2    def miceAndCheese(self, reward1: List[int], reward2: List[int], k: int) -> int:
3        # Get the total number of cheese pieces
4        n = len(reward1)
5      
6        # Sort indices based on the difference (reward1[i] - reward2[i]) in descending order
7        # This prioritizes cheese pieces where mouse 1 gets relatively more reward compared to mouse 2
8        indices = sorted(range(n), key=lambda i: reward1[i] - reward2[i], reverse=True)
9      
10        # Mouse 1 eats the first k cheese pieces (those with highest relative advantage)
11        # Mouse 2 eats the remaining cheese pieces
12        total_reward = sum(reward1[i] for i in indices[:k]) + sum(reward2[i] for i in indices[k:])
13      
14        return total_reward
15

1class Solution {
2    public int miceAndCheese(int[] reward1, int[] reward2, int k) {
3        int n = reward1.length;
4      
5        // Create an array of indices to track original positions
6        Integer[] indices = new Integer[n];
7        for (int i = 0; i < n; i++) {
8            indices[i] = i;
9        }
10      
11        // Sort indices based on the difference between rewards
12        // We want to maximize the gain from choosing reward1 over reward2
13        // Sort in descending order of (reward1[i] - reward2[i])
14        Arrays.sort(indices, new Comparator<Integer>() {
15            @Override
16            public int compare(Integer i, Integer j) {
17                int differenceI = reward1[i] - reward2[i];
18                int differenceJ = reward1[j] - reward2[j];
19                return differenceJ - differenceI; // Descending order
20            }
21        });
22      
23        int totalReward = 0;
24      
25        // First mouse eats k pieces of cheese from reward1
26        // Choose the k pieces where reward1 gives maximum advantage over reward2
27        for (int i = 0; i < k; i++) {
28            totalReward += reward1[indices[i]];
29        }
30      
31        // Second mouse eats the remaining (n - k) pieces from reward2
32        for (int i = k; i < n; i++) {
33            totalReward += reward2[indices[i]];
34        }
35      
36        return totalReward;
37    }
38}
39

1class Solution {
2public:
3    int miceAndCheese(vector<int>& reward1, vector<int>& reward2, int k) {
4        int n = reward1.size();
5      
6        // Create an index array to track original positions
7        vector<int> indices(n);
8        iota(indices.begin(), indices.end(), 0);
9      
10        // Sort indices based on the difference between rewards
11        // We want to maximize (reward1[i] - reward2[i]) for the first mouse
12        // So we sort in descending order of this difference
13        sort(indices.begin(), indices.end(), [&](int i, int j) {
14            return (reward1[i] - reward2[i]) > (reward1[j] - reward2[j]);
15        });
16      
17        int totalReward = 0;
18      
19        // First mouse eats k pieces of cheese with highest reward difference
20        for (int i = 0; i < k; ++i) {
21            totalReward += reward1[indices[i]];
22        }
23      
24        // Second mouse eats the remaining (n - k) pieces of cheese
25        for (int i = k; i < n; ++i) {
26            totalReward += reward2[indices[i]];
27        }
28      
29        return totalReward;
30    }
31};
32

1/**
2 * Calculates maximum reward when first mouse eats k pieces of cheese
3 * and second mouse eats the remaining pieces
4 * @param reward1 - Array of rewards if first mouse eats each piece
5 * @param reward2 - Array of rewards if second mouse eats each piece  
6 * @param k - Number of pieces the first mouse must eat
7 * @returns Maximum total reward possible
8 */
9function miceAndCheese(reward1: number[], reward2: number[], k: number): number {
10    const n: number = reward1.length;
11  
12    // Create array of indices [0, 1, 2, ..., n-1]
13    const indices: number[] = Array.from({ length: n }, (_, i) => i);
14  
15    // Sort indices by the difference in rewards (reward1 - reward2) in descending order
16    // Higher difference means first mouse gains more relative benefit from eating that piece
17    indices.sort((i: number, j: number) => {
18        const differencei: number = reward1[i] - reward2[i];
19        const differencej: number = reward1[j] - reward2[j];
20        return differencej - differencei;
21    });
22  
23    let totalReward: number = 0;
24  
25    // First mouse eats k pieces with highest relative benefit
26    for (let i = 0; i < k; ++i) {
27        totalReward += reward1[indices[i]];
28    }
29  
30    // Second mouse eats remaining pieces
31    for (let i = k; i < n; ++i) {
32        totalReward += reward2[indices[i]];
33    }
34  
35    return totalReward;
36}
37

Time and Space Complexity

Time Complexity: O(n × log n)

The time complexity is dominated by the sorting operation. The sorted() function with a custom key sorts the indices based on the difference reward1[i] - reward2[i], which requires O(n × log n) time. The subsequent operations - creating the range of indices O(n), computing the differences for sorting O(n), and the two sum operations O(k) and O(n-k) - are all linear and don't affect the overall complexity.

Space Complexity: O(n)

The space complexity comes from storing the sorted list of indices idx, which contains all n indices from 0 to n-1. This requires O(n) additional space. The sorting algorithm itself may use additional space (typically O(n) for Python's Timsort), but the dominant factor remains the storage of the index array.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall 1: Misunderstanding the Greedy Logic

The Problem: Many developers initially think they should simply pick the k highest values from reward1 for the first mouse and give the rest to the second mouse. This approach ignores the opportunity cost - what the second mouse could have gained from those same pieces.

Why it's wrong: Consider reward1 = [1, 1], reward2 = [1, 1000], and k = 1. Picking the highest from reward1 would give cheese 0 or 1 to mouse 1 (getting 1 point) and the other to mouse 2. But if mouse 2 gets cheese 1, we get 1000 points. The total would be 1 + 1 = 2 or 1 + 1000 = 1001. However, the optimal solution is to give cheese 0 to mouse 1 and cheese 1 to mouse 2, yielding 1 + 1000 = 1001 points.

The Solution: Always consider the relative advantage by calculating reward1[i] - reward2[i]. This difference tells us the net gain (or loss) of assigning cheese i to mouse 1 instead of mouse 2.

Pitfall 2: Incorrect Difference Calculation

The Problem: Some might calculate the difference as reward2[i] - reward1[i] instead of reward1[i] - reward2[i], or forget to reverse the sort order.

Why it's wrong: If you calculate reward2[i] - reward1[i] and sort in descending order, you'd be prioritizing cheese pieces where mouse 2 has the greatest advantage, which is the opposite of what we want when selecting cheese for mouse 1.

The Solution: Remember that we're selecting cheese for mouse 1, so we want reward1[i] - reward2[i] sorted in descending order. A positive difference means mouse 1 has an advantage for that cheese, while a negative difference means mouse 2 has an advantage.

Pitfall 3: Edge Case Handling

The Problem: Not considering edge cases like k = 0 or k = n.

Why it's wrong:

When k = 0, mouse 1 eats nothing, and all cheese goes to mouse 2
When k = n, mouse 1 eats everything, and mouse 2 gets nothing

The Solution: The current implementation handles these cases correctly:

When k = 0: indices[:0] is empty, so we sum all reward2 values
When k = n: indices[n:] is empty, so we sum all reward1 values

Always test these boundary conditions to ensure your solution works correctly.

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

How would you design a stack which has a function min that returns the minimum element in the stack, in addition to push and pop? All push, pop, min should have running time O(1).

Add another stack to store the minimums

Add a heap to store the minimums

Add a variable to track the minimum

Add a queue to store the minumums

I don't know