CodeIQ: バイナリ・カウント

CodeIQに公開された川添さんの問題「バイナリ・カウント」を解きました.
挑戦者求む!数学の問題をプログラミングで解こう!二進数の1の個数を数えて下さい。 by @riverplus Kawazoe│CodeIQ

問題

10進数表記の整数nが与えられる.次で定義する F(n) を求めよ.
   F(n) = \sum_{k = 1}^{n} (kを2進表記したときの1が立っているビット数)

整数nに対してn(10)を10進数表記、n(2)を2進数表記とします.

(例)n = 5 の場合

  k = 1(10), 001(2), 1が立っているビット数 = 1
  k = 2(10), 010(2), 1が立っているビット数 = 1
  k = 3(10), 011(2), 1が立っているビット数 = 2
  k = 4(10), 100(2), 1が立っているビット数 = 1
  k = 5(10), 101(2), 1が立っているビット数 = 2
   \therefore F(5) = 1 + 1 + 2 + 1 + 2 = 7

アルゴリズム

F(12)を例に規則性を見てみます.
f:id:tatanaideyo:20150430225327p:plain
i桁目( 0 \le i \le 3)の列を見てみると、整数  2^i から周期  2^i で連続して1と0のビットが立っているのが分かります.したがって、各々の桁に対して1が立っている整数の数を数えます.
n = 19のi = 2桁目に対して1が立っている整数の数を数えます.
f:id:tatanaideyo:20150501003308p:plain
整数kのidxが偶数ならidx桁目のビットは1、奇数なら0となります.
nのidxが奇数の場合は  \frac{idx + 1}{2} \times 2^{digit} 個、nのidxが偶数の場合は  \frac{idx}{2} \times 2^{digit} + n - (idx \times 2^{digit} + 2^{digit} - 1) 個の整数のidx桁目に1が立っています.
これを、 (int)\log_{10} n + 1 桁分に対して総和を求めると答えになります.

 O(1) で出来ないものかな)

時間計算量

整数nの各桁に対して  O(1) で処理しているので全体で  O(\log n) です.

ソースコード

n = 10000000000
ans = 0
digit = 0

while 2**digit <= n do
  offset = (n - 2**digit)
  idx = offset / (2**digit)

  if idx % 2 == 1 then
    ans += ((idx + 1) / 2) * 2**digit
  else
    ans += (idx / 2) * 2**digit
    ans += n - (idx * 2**digit + 2**digit - 1)
  end
  
  digit += 1
end

puts ans