Pythonで0からnまでの全数値のセットビット総数をカウントするプログラム

問題の概要

ある整数 num が与えられたとき、0 ≤ i ≤ num の範囲に含まれる各整数 i について、その2進数表現における「1」の個数(セットビット数)を求めます。

例えば、num が 5 の場合、対象となる数値は [0, 1, 2, 3, 4, 5] です。それぞれを2進数で表すと次のようになります。

• 0 → 0 → セットビット数 0
• 1 → 1 → セットビット数 1
• 2 → 10 → セットビット数 1
• 3 → 11 → セットビット数 2
• 4 → 100 → セットビット数 1
• 5 → 101 → セットビット数 2

したがって、各数値のセットビット数は [0, 1, 1, 2, 1, 2] となり、その合計値である 7 を返します。

アルゴリズム

この問題は、ビット演算を活用した動的計画法(DP)の考え方で効率的に解けます。手順は以下の通りです。

• res := num + 1 個の 0 を要素にもつ配列を用意する

• offset := 0 と初期化する

• i を 1 から num まで順に処理する:

  • i & (i − 1) == 0 の場合(つまり i が2のべき乗の場合)、res[i] := 1 とし、offset := 0 にリセットする

  • それ以外の場合は、offset を 1 増やし、res[i] := 1 + res[offset] とする

• 最後に、res の全要素の合計を返す

なぜこの方法が機能するのか

ポイントは条件式 i & (i − 1) です。これは i が2のべき乗(1, 2, 4, 8, …)のときだけ結果が 0 になるビット演算です。2のべき乗の2進数表現は最上位に「1」を1つだけ持つためです。

i が2のべき乗でない場合、i は「直前の2のべき乗 + offset」という形で表せます。このとき、i のセットビット数は「先頭ビット分の 1」と「offset のセットビット数」の和になります。つまり、すでに計算済みの小さい数値の結果を再利用できるため、各数値を実質 O(1) で処理できます。

実装例

以下がPythonでの実装コードです。

class Solution:
    def countBits(self, num):
        result = [0] * (num + 1)
        offset = 0
        for i in range(1, num + 1):
            if i & (i - 1) == 0:
                result[i] = 1
                offset = 0
            else:
                offset += 1
                result[i] = 1 + result[offset]
        return sum(result)

ob1 = Solution()
print(ob1.countBits(5))

入力

5

出力

7

計算量

このアルゴリズムは 0 から num までの各数値を一度ずつ処理するため、時間計算量は O(num)、結果を保存する配列が必要なため空間計算量も O(num) です。各数値について2進数へ変換しながらビットを数える素朴な方法(O(num log num))と比べて効率的です。

別のアプローチ:よりシンプルな漸化式

参考までに、さらに簡潔に書ける方法もあります。result[i] = result[i >> 1] + (i & 1) という関係を利用すると、「i のセットビット数 = i を1ビット右シフトした値のセットビット数 + i の最下位ビット」として計算できます。

class Solution:
    def countBits(self, num):
        result = [0] * (num + 1)
        for i in range(1, num + 1):
            result[i] = result[i >> 1] + (i & 1)
        return sum(result)