Pythonで上限値以下の要素の中からXORが最大となる数を見つけるプログラム

問題概要

数値のリスト nums と、各クエリが [x, limit] の形式を持つクエリリスト queries が与えられます。各クエリに対して、nums の中から limit 以下の要素 e を探し、e XOR x が最大になるような e を求めます。条件を満たす要素が存在しない場合は -1 を返します。

例えば、nums = [3, 5, 9]、queries = [[4, 6], [2, 0]] の場合、出力は [3, -1] となります。最初のクエリでは、6以下の要素として 3 と 5 が候補になります。3 XOR 4 = 7、5 XOR 4 = 1 なので、より大きなXOR値を返す 3 が選ばれます。2番目のクエリでは、0以下の数値が nums に存在しないため、-1 を返します。

解決のアプローチ

この問題は「ビット単位のトライ（Trie）」というデータ構造を使うと効率的に解けます。各数値を32ビットの二進表現としてトライに挿入し、クエリ時には反転ビットを優先して辿ることで、XORを最大化する要素を高速に特定できます。

さらに、配列 A とクエリをそれぞれソートしておくことで、各クエリの limit 以下の要素だけを段階的にトライへ挿入でき、無駄な探索を避けられます。

具体的な手順は以下の通りです。

• trie を空のマップとして初期化する

• 関数 bits(i) を定義する：i の32ビット二進表現を返す

• 関数 insert(i) を定義する：

  • node := trie

  • bits(i) の各ビット c について、node := node.setdefault(c, {}) で子ノードを辿る

  最後に node[2] := i を設定する

• 関数 query(i) を定義する：

  • node := trie

  • bits(i) の各ビット c について、rc := c XOR 1（反転ビット）とし、node[rc] が存在すればそちらへ、なければ node[c] へ移動する

  最後に node[2] を返す

• メイン処理：

  • リスト A を昇順にソートする

  • B := 各クエリを (インデックス i, x, limit) のタプルに変換したリストを作り、limit を基準にソートする

  • (j, n, ans) := (0, Aのサイズ, クエリ数と同じ長さの -1 で埋めたリスト)

  • B の各 (i, x, limit) について：

    • j < n かつ A[j] <= limit の間、insert(A[j]) を実行し、j を 1 ずつ増やす

    • j が 0 でなければ、ans[i] := query(x)

  • ans を返す

なぜ反転ビットを優先するのか

XORを最大化するには、x の各ビットに対して反対のビット（0なら1、1なら0）を持つ数を選ぶのが理想的だからです。query() 内で rc = c XOR 1 の子ノードを先に試すのは、まさにこの性質を利用しています。

Pythonでの実装例

以下の実装を見ると、理解が深まるでしょう。

class Solution:
    def solve(self, A, queries):
        trie = {}
        def bits(i):
            return map(int, bin(i)[2:].zfill(32))
        def insert(i):
            node = trie
            for c in bits(i):
                node = node.setdefault(c, {})
            node[2] = i
        def query(i):
            node = trie
            for c in bits(i):
                rc = c ^ 1
                node = node.get(rc, node.get(c))
            return node[2]
        A.sort()
        B = sorted([(i, x, limit) for i, (x, limit) in enumerate(queries)], key=lambda x: x[2])
        j, n, ans = 0, len(A), [-1] * len(queries)
        for i, x, limit in B:
            while j < n and A[j] <= limit:
                insert(A[j])
                j += 1
            if j:
                ans[i] = query(x)
        return ans
ob = Solution()
nums = [3, 5, 9]
queries = [
    [4, 6],
    [2, 0]
]
print(ob.solve(nums, queries))

入力

[3, 5, 9], [[4, 6],[2, 0]]

出力

[3, -1]

まとめ

本手法では、数値の挿入とクエリ処理がそれぞれビット幅（32回）の走査で済むため、全体の計算量は O((N + Q) × 32) 程度に抑えられます。ソートによって limit の制約を前処理で解消している点も、このアルゴリズムの重要なポイントです。