Leetcode 0781 Rabbits in Forest

Problem statement

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

Solution

Note, that for if some of the rabbits says number i, than there is group of i+1 rabbits with one color. So, if we have 1, 2, ..., i+1 of them saying i, minimal number of rabbits for this group is i+1, if we have i+2, ..., 2i+2, minimal number of rabbits is 2i+2: there should be 2 groups. In general it can be written as (j+i)//(i+1)*(i+1) if we have j rabbits which said i.

Complexity

Finally, all this can be written in oneliner, using counters, time and space complexity is O(n).

Code

class Solution:
    def numRabbits(self, answers):
        return sum((j+i)//(i+1)*(i+1) for i,j in Counter(answers).items())