781. Rabbits in Forest

Posted 2021-01-03 learning-c

class Solution(object):
    def numRabbits(self, answers):
        """
        :type answers: List[int]
        :rtype: int
        """
        dict={}
        for x in answers:
            if x in dict.keys():
                dict[x]=dict[x]+1
            else:
                dict[x]=1
        cnt=0
        for k in dict.keys():
            v=dict[k]
            if v%(k+1)==0:
                cnt+=v/(k+1)*(k+1)
            else:
                cnt+=v/(k+1)*(k+1)+(k+1)
        return cnt