LeetCode 781: Rabbits in Forest

Problem Restatement

We are given an array answers.

Each value answers[i] is the answer from one rabbit.

If a rabbit says x, it means:

There are x other rabbits with the same color as me.

So that rabbit belongs to a color group of size:

x + 1

We need to return the minimum possible number of rabbits in the forest.

Some rabbits in the forest may not have answered the question.

Input and Output

Function shape:

class Solution:
    def numRabbits(self, answers: list[int]) -> int:
        ...

Examples

Example 1:

answers = [1, 1, 2]

The two rabbits that answered 1 could have the same color.

If a rabbit answers 1, its color group has size 2.

So the two 1 answers can be one group of size 2.

The rabbit that answered 2 belongs to a group of size 3.

Only one rabbit from that group answered, so there must be two more rabbits of that color in the forest.

Total:

2 rabbits from answer 1 group
3 rabbits from answer 2 group
= 5 rabbits

So the answer is:

5

Example 2:

answers = [10, 10, 10]

Each rabbit says there are 10 other rabbits with the same color.

So each such color group has size:

10 + 1 = 11

All three rabbits can belong to the same group of size 11.

So the answer is:

11

First Thought: Count Every Rabbit Directly

A simple but wrong idea is to say:

answer = len(answers)

This only counts rabbits that answered.

But the forest may contain rabbits that did not answer.

For example:

answers = [2]

One rabbit says there are 2 other rabbits with the same color.

So there must be at least 3 rabbits total.

The answer is not 1.

It is:

3

Problem With Naive Counting

Each answer gives information about a whole color group.

If a rabbit answers x, then its color group has exactly x + 1 rabbits.

So rabbits with the same answer may be grouped together, but only up to x + 1 rabbits per group.

For example:

answers = [1, 1, 1]

Each 1 means a color group of size 2.

Only two rabbits can fit into one such group.

The third rabbit must belong to another color group of size 2.

So the minimum total is:

2 + 2 = 4

not 3.

Key Insight

Group rabbits by their answer.

Suppose an answer value x appears count times.

Each group for answer x has size:

group_size = x + 1

At most group_size answering rabbits can fit into one group.

So the number of groups needed is:

ceil(count / group_size)

Each group contributes group_size rabbits to the total.

So the contribution is:

ceil(count / group_size) * group_size

In integer arithmetic:

groups = (count + group_size - 1) // group_size

Algorithm

1. Count how many times each answer appears.
2. For each answer x:
   1. Let group_size = x + 1.
   2. Let count be how many rabbits gave answer x.
   3. Compute how many groups are needed.
   4. Add groups * group_size to the answer.
3. Return the total.

Correctness

For any rabbit that answers x, its color group must contain exactly x + 1 rabbits.

Therefore, among rabbits that answered x, one color group can contain at most x + 1 of them. If more than x + 1 rabbits answered x, they cannot all have the same color. They must be split into multiple color groups.

For a fixed answer x, let count be the number of rabbits that gave that answer. Since each group can hold at most x + 1 answering rabbits, we need at least:

ceil(count / (x + 1))

groups.

Each such group has exactly x + 1 rabbits, including rabbits that may not have answered. So the minimum number of rabbits contributed by answer x is:

ceil(count / (x + 1)) * (x + 1)

The algorithm computes exactly this value for every distinct answer.

Rabbits with different answers cannot belong to the same color group, because all rabbits of one color would give the same answer. Therefore, the contributions from different answer values are independent and can be added.

Thus, the algorithm returns the minimum possible number of rabbits in the forest.

Complexity

Here, n is the length of answers.

Implementation

from collections import Counter

class Solution:
    def numRabbits(self, answers: list[int]) -> int:
        counts = Counter(answers)
        total = 0

        for x, count in counts.items():
            group_size = x + 1
            groups = (count + group_size - 1) // group_size
            total += groups * group_size

        return total

Code Explanation

We count equal answers:

counts = Counter(answers)

For each answer x, the color group size is:

group_size = x + 1

If x = 0, the group size is 1. That means each such rabbit is alone in its color group.

If x = 2, the group size is 3.

We compute the number of groups using integer ceiling division:

groups = (count + group_size - 1) // group_size

Then we add all rabbits from those groups:

total += groups * group_size

Finally, return the minimum total:

return total

Testing

def run_tests():
    s = Solution()

    assert s.numRabbits([1, 1, 2]) == 5
    assert s.numRabbits([10, 10, 10]) == 11
    assert s.numRabbits([]) == 0
    assert s.numRabbits([0, 0, 0]) == 3
    assert s.numRabbits([1, 1, 1]) == 4
    assert s.numRabbits([2, 2, 2]) == 3
    assert s.numRabbits([2, 2, 2, 2]) == 6

    print("all tests passed")

run_tests()

Test coverage: