s0338_counting_bits

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:

• 0 <= n <= 105

Follow up:

• It is very easy to come up with a solution with a runtime of O(n log n). Can you do it in linear time O(n) and possibly in a single pass?
• Can you do it without using any built-in function (i.e., like __builtin_popcount in C++)?

Solution

defmodule Solution do
  @spec count_bits(n :: integer) :: [integer]
  def count_bits(n), do: do_count_bits(1, n, Map.new(0..div(n,2), &({&1,0})), [0])

  defp do_count_bits(i, n, map, acc) when i > n, do: Enum.reverse(acc)
  defp do_count_bits(i, n, map, acc) do
    i_bits = map[div(i,2)] + rem(i,2)
    new_map = if i > div(n,2), do: map, else: Map.put(map, i, i_bits)
    do_count_bits(i + 1, n, new_map, [i_bits | acc])
  end
end