LeetCode 338. Counting Bits

counting-bits.java

/**
 * https://leetcode.com/problems/counting-bits/
 *
 * O(n) runtime, O(1) space
 *
 * use DP to solve this question
 *
 * assume k is a multiple of two
 * then for k < i < 2k, f(i) = 1 + f(i - k)
 * where f(x) means the number of 1's in x's binary representation
 */

public class Solution {
    public int[] countBits(int num) {
        if (num == 0) {
            return new int[] { 0 };
        }
        int[] res = new int[num+1];
        res[1] = 1; // initialization
        for (int i = 2, weight = 1; i <= num; i++) {
            if (i < 2 * weight) {
                res[i] = 1 + res[i-weight];
            } else {
                weight *= 2;
                res[i] = 1;
            }
        }
        return res;
    }
}