223. Counting Bits

Dropping the last bit of i gives i // 2, whose bit count we already computed. Add back the bit we dropped (i & 1).

Key relation

The binary form of i >> 1 is i with its last bit removed. So the number of set bits in i equals the set bits in i >> 1 plus the last bit itself:

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

Since i >> 1 < i, its answer is already filled — a clean bottom-up DP.

Walk-through

1. Allocate ans of length n + 1, with ans[0] = 0.
2. For i from 1 to n: ans[i] = ans[i >> 1] + (i & 1).
3. Return ans.

Worked example (n = 5)

Result: [0,1,1,2,1,2].

Complexity

• Time: O(n) — one O(1) transition per value.
• Space: O(1) extra (ignoring the output array).

class Solution {
    public int[] countBits(int n) {
        int[] ans = new int[n + 1];
        for (int i = 1; i <= n; i++) {
            ans[i] = ans[i >> 1] + (i & 1);
        }
        return ans;
    }
}