1
# Runtime: 54 ms (Top 35.1%) | Memory: 16.38 MB (Top 76.8%)
2

3

4
class Solution:
5
    def numRabbits(self, answers: List[int]) -> int:
6
        counts = {}
7
        count = 0
8

9
        for answer in answers:
10
            if answer == 0:
11
                # This must be the only rabbit of its color.
12
                count += 1
13
            elif answer not in counts or counts[answer] == 0:
14
                # This is the first time the color appears.
15
                counts[answer] = 1
16
                # Add all rabbits having this new color.
17
                count += answer + 1
18
            else:
19
                # Add one to how many times answer occurred.
20
                counts[answer] += 1
21
                if counts[answer] > answer:
22
                    # If n+1 rabbits have said n,
23
                    # this color group is complete.
24
                    counts[answer] = 0
25

26
        return count