[LeetCode] 338. Counting Bits

338. Counting Bits

Given a non negative integer number num. For every numbers i in the range 0 ≤ i ≤ num calculate the number of 1's in their binary representation and return them as an array.
Example:
For num = 5 you should return [0,1,1,2,1,2].

Solution:

class Solution:
    def countBits(self, num):
        """
        :type num: int
        :rtype: List[int]
        """
        if num == 0:
            return [0]
        count = [0] * (num + 1)
        a = 1
        for i in range(1, num+1):
            if i == a:
                count[i] = 1
                a *= 2
            else:
                count[i] = count[a//2] + count[i - a//2]
        return count