Counting Bits — Leetcode 338

Original Leetcode problem: https://leetcode.com/problems/counting-bits/

In this problem, we’re given an n and asked to return an array of size n + 1 that contains the number of 1 bits of every number 0 ≤ i ≤ n.

Before studying the Leetcode O(n) solution, I came up with the intuitive O(nlogn) solution:

class Solution {
    public int[] countBits(int n) {        
        int[] result = new int[n + 1];
        for (int i = 0; i <= n; i++) {
            String binary = Integer.toBinaryString(i);
            int count = 0;
            for (char ch : binary.toCharArray()) {
                if (ch == '1') {
                    count++;
                }
            }
            result[i] = count;
        }
        return result;
    }
}

I used this to plot the following chart in Excel to help me find a pattern for the O(n) solution suggested in the Leetcode follow up:

I generated the array from my O(nlogn) solution for n = 63. Combined with studying the Leetcode solution article, I was able to see the pattern easily from the graph: for every range between powers of 2 (e.g. range [powerOfTwo, nextPowerOfTwo — 1] such as [2, 3], [4, 7], [8, 15], etc) the count of 1’s for the range was equivalent to 1 + count[0,powerOfTwo — 1]. For example, to get the ones count for [2,3], we notice from the graph that for [0, 1] the count is 0, 1 respectively and for [2, 3] it is 1, 2 respectively. Similarly, for [4, 7] it is 0+1=1, 1+1=2, 1+1=2, 2+1=3 respectively (deriving from the counts for range [0, 3]). We can extrapolate this pattern and get the following O(n) solution:

class Solution {
    public int[] countBits(int n) {
        int[] result = new int[n + 1];
        result[0] = 0;
        if (n == 0) {
            return result;
        }
        result[1] = 1;
        int i = 2;
        int powerOfTwo = 4;
        while (i <= n) {
            while (i <= n && i < powerOfTwo) {
                result[i] = result[i - (powerOfTwo >> 1)] + 1;
                i++;
            }
            powerOfTwo = powerOfTwo << 1;
        }
        return result;
    }
}