338_counting_bits.cpp

/*
LeetCode 338 - Counting Bits
Difficulty: Easy

Problem:
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]

Example 2:
    Input: n = 5
    Output: [0,1,1,2,1,2]

Constraints:
    0 <= n <= 10^5

Time Complexity: O(n)
Space Complexity: O(n)  (for the output array)
*/

#include <iostream>
#include <vector>

using namespace std;

class Solution {
public:
    vector<int> countBits(int n) {
        vector<int> dp(n + 1, 0);
        for (int i = 1; i <= n; i++) {
            dp[i] = dp[i >> 1] + (i & 1);
        }
        return dp;
    }
};

int main() {
    Solution solution;

    vector<int> result1 = solution.countBits(2);
    cout << "[";
    for (int i = 0; i < result1.size(); i++) {
        cout << result1[i];
        if (i < result1.size() - 1) cout << ", ";
    }
    cout << "]" << endl;  // [0, 1, 1]

    vector<int> result2 = solution.countBits(5);
    cout << "[";
    for (int i = 0; i < result2.size(); i++) {
        cout << result2[i];
        if (i < result2.size() - 1) cout << ", ";
    }
    cout << "]" << endl;  // [0, 1, 1, 2, 1, 2]

    return 0;
}