[LeetCode] 781. Rabbits in Forest

 Problem : https://leetcode.com/problems/rabbits-in-forest/

Remember the last rabbit answered the same values. Then count rabbits must be in different colors.

class Solution:
    def numRabbits(self, answers: List[int]) -> int:
        
        total = 0
        seen = defaultdict(int)
        
        for n in answers:
            if n == 0:
                # encounter a single rabbit of this color
                total += 1
            else:
                if seen[n+1]:
                    # encounter rabbit with same color
                    seen[n+1] -= 1
                else:
                    # encounter rabbit with different color
                    seen[n+1] += n
                    total += n+1
                    
        return total

Another math solution:

Firstly count number of rabbits anwsered the same values. For example, M rabbits answered N. Then there must be M // (N+1) group of rabbits with same color. Each group has N+1 rabbits. Additionally, there must be another group of N+1 rabbits if M % (N+1) > 0.

class Solution:
    def numRabbits(self, answers: List[int]) -> int:
        
        total = 0
        count = Counter(answers)
        
        for n, m in count.items():
            total += m // (n+1) * (n+1)
            total += n + 1 if m % (n+1) else 0
                    
        return total