781. Rabbits in Forest

Posted 2021-01-05 ffeng0312

class Solution(object):
    def numRabbits(self, answers):
        """
        :type answers: List[int]
        :rtype: int
        """
        cnt = {}
        for answer in answers:
            if answer not in cnt:
                cnt[answer] = 1
            else:
                cnt[answer] += 1

        s = 0
        for k, v in cnt.iteritems():
            if v <= k+1:
                s += k+1
            elif k > 0:
                rest = v%(k+1)
                if rest <>0:
                    s += v - rest
                    s += k+1
                else:
                    s += v
            elif k == 0:
                s += v

        return s