Counting Bits

Difficulty: Easy Source: NeetCode

Problem

Given an integer n, return an array ans of length n + 1 such that for each i (0 <= i <= n), ans[i] is the number of 1’s in the binary representation of i.

Example 1: Input: n = 2 Output: [0, 1, 1] Explanation: 0 → 0, 1 → 1, 2 → 10

Example 2: Input: n = 5 Output: [0, 1, 1, 2, 1, 2] Explanation: 0 → 0, 1 → 1, 2 → 10, 3 → 11, 4 → 100, 5 → 101

Constraints:

• 0 <= n <= 10⁵

Follow up: Can you do it in O(n) without any built-in functions?

Prerequisites

Before attempting this problem, you should be comfortable with:

• Right shift (>>) — i >> 1 is the same as i // 2, removing the last bit
• Bitwise AND (&) — i & 1 extracts the last bit (0 or 1)
• Dynamic programming — building answers for larger inputs from previously computed smaller inputs

1. Brute Force

Intuition

Count the set bits for every number from 0 to n independently. We already know how to count bits in a single number (see Number of 1 Bits). Just call that logic for each number and collect the results. It’s correct but wasteful — we’re recomputing work we’ve already done for smaller numbers.

Algorithm

1. Initialize ans = [].
2. For each i from 0 to n (inclusive):
   • Count the 1-bits in i using the bit-shifting method.
   • Append the count to ans.
3. Return ans.

Solution

def countBits(n):
    def count_ones(x):
        count = 0
        while x > 0:
            count += x & 1
            x >>= 1
        return count

    return [count_ones(i) for i in range(n + 1)]


print(countBits(2))  # [0, 1, 1]
print(countBits(5))  # [0, 1, 1, 2, 1, 2]
print(countBits(0))  # [0]

Complexity

• Time: O(n log n) — for each of the n numbers, we inspect up to log n bits
• Space: O(n) — for the output array

2. Dynamic Programming with Bit Trick

Intuition

Here’s the key insight: every number i is just i >> 1 (i.e., i // 2) with possibly one extra bit appended on the right.

• i >> 1 removes the last bit — we already know how many 1s are in i >> 1 from our earlier computation.
• The last bit is either 0 or 1, and we can read it with i & 1.

So: dp[i] = dp[i >> 1] + (i & 1)

Let’s see this for a few numbers:

i = 0  (0)     →  dp[0 >> 1] + (0 & 1) = dp[0] + 0 = 0 + 0 = 0
i = 1  (1)     →  dp[1 >> 1] + (1 & 1) = dp[0] + 1 = 0 + 1 = 1
i = 2  (10)    →  dp[2 >> 1] + (2 & 1) = dp[1] + 0 = 1 + 0 = 1
i = 3  (11)    →  dp[3 >> 1] + (3 & 1) = dp[1] + 1 = 1 + 1 = 2
i = 4  (100)   →  dp[4 >> 1] + (4 & 1) = dp[2] + 0 = 1 + 0 = 1
i = 5  (101)   →  dp[5 >> 1] + (5 & 1) = dp[2] + 1 = 1 + 1 = 2

Notice i >> 1 is always less than i, so the value we need is always already computed by the time we get to i.

Algorithm

1. Initialize dp = [0] * (n + 1).
2. For i from 1 to n:
   • Set dp[i] = dp[i >> 1] + (i & 1).
3. Return dp.

Solution

def countBits(n):
    dp = [0] * (n + 1)
    for i in range(1, n + 1):
        dp[i] = dp[i >> 1] + (i & 1)
    return dp


print(countBits(2))  # [0, 1, 1]
print(countBits(5))  # [0, 1, 1, 2, 1, 2]
print(countBits(0))  # [0]

Complexity

• Time: O(n)
• Space: O(n) — for the output array (no extra space beyond the answer)

Common Pitfalls

Starting the loop at 0. dp[0] is 0 by initialization and we’d compute dp[0 >> 1] + (0 & 1) = dp[0] + 0 = 0 anyway — it’s harmless, but starting at 1 is cleaner and avoids any confusion.

Confusing i >> 1 with i - 1. Right shift by 1 is integer division by 2, not subtraction. 5 >> 1 = 2, not 4.

Using // instead of >>. Both work correctly for non-negative integers, but >> makes the bitwise intent explicit and is idiomatic for this type of problem.