LeetCode-in-All

338. Counting Bits

Easy

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:

Follow up:

Solution

#include <stdio.h>
#include <stdlib.h>

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
int* countBits(int num, int* returnSize) {
    // Allocate memory for the result array and set the return size
    int* result = (int*)malloc((num + 1) * sizeof(int));
    *returnSize = num + 1;

    int borderPos = 1;
    int incrPos = 1;

    result[0] = 0; // Base case: 0 has zero 1-bits

    for (int i = 1; i <= num; i++) {
        // When we reach a power of 2, reset borderPos and incrPos
        if (incrPos == borderPos) {
            result[i] = 1;  // Powers of 2 have exactly one 1-bit
            incrPos = 1;
            borderPos = i;
        } else {
            result[i] = 1 + result[incrPos++];
        }
    }

    return result;
}