Pythonでビットを数える：0〜numまでの各整数の2進数における「1」の個数をO(n)で求める方法

問題概要

非負整数 num が与えられたとき、0 ≤ i ≤ num を満たす各整数 i について、その2進表現に含まれる「1」の個数を計算し、結果をリストとして返すことを考えます。

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

解法のアプローチ（動的計画法）

この問題は、すでに計算した結果を再利用する動的計画法（DP）の考え方を使うことで、各数値を1度ずつ処理するだけで全体を O(n) の時間計算量で解けます。

ポイントは次の2つです。

• 2のべき乗の判定： i & (i - 1) == 0 が成り立つとき、i は2のべき乗です。2のべき乗の2進表現には「1」が1つしか含まれないため、result[i] = 1 とできます。
• オフセットの利用： 直前の2のべき乗を基準とした相対位置 offset を管理します。i が2のべき乗でない場合、i の2進表現は「offset の2進表現の先頭に1を付け加えたもの」と一致するため、result[i] = 1 + result[offset] という漸化式が成り立ちます。

アルゴリズムの手順

• 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 を返す

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 result

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

入力

6

出力

[0, 1, 1, 2, 1, 2, 2]

なお、Pythonではビット演算子 & は比較演算子よりも優先順位が高いため、i & i - 1 == 0 は (i & (i - 1)) == 0 と解釈され、括弧がなくても意図どおりに動作します。

別解との比較

より簡潔に書きたい場合は、組み込み関数を使った次の方法もあります。

def count_bits(num):
    return [bin(i).count('1') for i in range(num + 1)]

この方法は直感的で読みやすい一方、各数値ごとに2進文字列を生成するため、入力が大きくなるとDPを使った手法の方が効率的です。また、result[i] = result[i >> 1] + (i & 1) という漸化式を使えば、オフセット変数なしで同様のDPを実装することも可能です。

計算量

• 時間計算量： O(n) — 各数値を1度ずつ処理するだけです
• 空間計算量： O(n) — 結果を格納する配列が必要です