Counting Bits (LeetCode 338)

The problem

Given an integer n, return a list ans of length n + 1 where ans[i] is the number of 1 bits in the binary form of i, for every i from 0 to n.

Input:  5
Output: [0, 1, 1, 2, 1, 2]

0 to 5 in binary are 0, 1, 10, 11, 100, 101, which have 0, 1, 1, 2, 1, 2 ones.

How to think about it

Start with brute force. Run Number of 1 Bits on every i. Each call costs up to O(log i), so the total is O(n log n). The follow-up asks for O(n), one pass, and no built-in popcount.

The key insight: i is i >> 1 with one extra bit on the end. Shifting right by one drops the last bit. So i has exactly the same 1 bits as i >> 1, plus one more if its last bit is 1. That last bit is i & 1.

bits[i] = bits[i >> 1] + (i & 1)

It is a DP because i >> 1 < i. By the time you compute bits[i], the answer for i >> 1 is already in the table. One addition per number gives O(n).

Which pattern is this? Bit manipulation feeding a simple dynamic programme. The tell is “answer for every number from 0 to n”, where each answer can be built from a smaller number’s answer. A second recurrence also works: bits[i] = bits[i & (i - 1)] + 1, which drops the lowest set bit instead.

Approach

1. Make a list dp of n + 1 zeros; dp[0] = 0 is the base case.
2. For i from 1 to n, set dp[i] = dp[i >> 1] + (i & 1).
3. Return dp.

Solution

def count_bits(n):
    dp = [0] * (n + 1)
    for i in range(1, n + 1):
        dp[i] = dp[i >> 1] + (i & 1)   # bits of i without its last bit, plus that bit
    return dp
# count_bits(5) -> [0,1,1,2,1,2]

Dry run

On n = 5:

The result is [0, 1, 1, 2, 1, 2].

Complexity

• Time: O(n). One constant-time update per number.
• Space: O(n) for the output list, which the problem requires anyway. No other extra space.

Edge cases and pitfalls

• n = 0 returns [0]; the loop does not run.
• Off by one. The list needs n + 1 entries, because 0 is included.
• Using i // 2 and i % 2 is the same thing; the shift and mask just state the bit idea directly.
• Calling popcount per number is the O(n log n) answer interviewers want you to improve on.

Related problems

• Number of 1 Bits (LeetCode 191): the single-number version, with n & (n - 1).
• Climbing Stairs (LeetCode 70): another one-dimensional DP built from smaller answers.
• Reverse Bits (LeetCode 190): more practice with >> and & 1.