【Python】ビット単位ORがKと等しくなるN個の異なる数を見つける方法

2つの整数 N と K が与えられたとき、それらのビット単位のOR（論理和）を計算すると結果がちょうど K と等しくなるような、N個の互いに異なる値を見つけることを考えます。条件を満たす組み合わせが存在しない場合は -1 を返します。

たとえば、入力が N = 4、K = 6 の場合、出力は [6, 0, 1, 2] となります。実際に確認すると、6 OR 0 OR 1 OR 2 = 6 となり、条件を満たしていることがわかります。

解法のアプローチ

この問題は、次の手順で解くことができます。

• MAX := 32 — 扱うビット幅を32ビットとします。
• visited: サイズ MAX のリストを作成し、すべて False で初期化します。使用済みのビット位置を管理するために使われます。
• res: 結果を格納するための空のリストです。
• 関数 add(num) を定義します。
  • point := 0、value := 0 で初期化します。
  • i を 0 から MAX-1 まで繰り返します。
    • visited[i] が True の場合は、次の反復へスキップします。
    • そうでない場合、num の最下位ビットが 1 であれば、value に 2^i を加算します。
    • num を 2 で割り（整数部分のみ）、次のビット判定へ進みます。
  • 最後に value を res の末尾に追加します。

続いて、メイン処理では以下を行います。

• pow2: 2^0 から 2^31 までの2のべき乗を格納した配列を作成します。
• まず K 自身を res に追加します。
• cnt_k: K を2進数で表したときのセットビット（1になっているビット）の数を求めます。
• もし pow2[cnt_k] < n であれば、N 個の異なる数を用意できないため -1 を返します。
• count := 0 とし、i を 0 から pow2[cnt_k] - 1 まで順に add(i) を呼び出していきます。
• count が n に達した時点でループを抜けます。
• 最後に res を返します。

なぜこの方法が成立するのか

K のセットビット数が cnt_k であるとき、0 から 2^cnt_k − 1 までの非負整数は、K の最上位ビットより上のビットを一切持ちません。そのため、これらの数をどれだけ K と OR しても、結果は常に K のままになります。この性質を利用すれば、答えとなる数列を簡単に構築できます。

一方で、候補となる数の総数（K 自身を含めて 2^cnt_k 個）が N 未満の場合は、重複なしに N 個を選ぶことが不可能なため、-1 を返す必要があります。

実装例

理解を深めるために、以下のPythonコードを見てみましょう。

MAX = 32
visited = [False for i in range(MAX)]
res = []

def set_bit_count(n):
    if (n == 0):
        return 0
    else:
        return (n & 1) + set_bit_count(n >> 1)

def add(num):
    point = 0
    value = 0
    for i in range(MAX):
        if (visited[i]):
            continue
        else:
            if (num & 1):
                value += (1 << i)
            num = num//2
    res.append(value)

def solve(n, k):
    pow2 = [2**i for i in range(MAX)]
    res.append(k)
    cnt_k = set_bit_count(k)
    if (pow2[cnt_k] < n):
        return -1
    count = 0
    for i in range(pow2[cnt_k] - 1):
        add(i)
        count += 1
        if (count == n):
            break
    return res

n = 4
k = 6
print(solve(n, k))

入力

4, 6

出力

[6, 0, 1, 2]