Pythonで配列の要素から最大XORを求めるプログラムの作成方法

非負の整数だけを含む配列 nums と、複数のクエリをまとめた配列 queries が与えられるとします。queries[i] はペア (xi, mi) を表しています。i 番目のクエリに対する答えは、mi 以下の nums の要素の中から選んだ任意の要素と xi のビット単位 XOR の最大値です。もし nums のすべての要素が mi より大きければ、その答えは -1 となります。

最終的には、queries と同じ長さの配列 answer を返す必要があります。answer[i] には i 番目のクエリの答えを格納してください。

入力例と出力例

例として、入力が nums = [0,1,2,3,4]、queries = [[3,1],[1,3],[5,6]] の場合を考えてみましょう。このときの出力は [3,3,7] になります。

• クエリ [3,1]: 1 以下の要素は 0 と 1 です。0 XOR 3 = 3、1 XOR 3 = 2 なので、大きい方の 3 が答えになります。
• クエリ [1,3]: 3 以下の要素との XOR の最大値は 1 XOR 2 = 3 です。
• クエリ [5,6]: 6 以下の要素との XOR の最大値は 5 XOR 2 = 7 です。

解法の考え方

この問題は、クエリを制限値(limit)の昇順に並べ替えてオフラインで一括処理し、上位ビットから順に答えを確定させていく貪欲法で効率よく解けます。具体的な手順は次の通りです。

• m := nums のサイズ
• n := queries のサイズ
• queries を、インデックス i とペア (x, limit) からなるトリプル (i, x, limit) に変換する
• limit を基準に queries をソートする
• nums を昇順にソートする
• res := サイズ n の配列を作成し、すべて 0 で初期化する
• k を 31 から 0 まで 1 ずつ減らしながら、以下を繰り返す
  • prefixes := 新しい空の集合(set)を作る
  • j := 0
  • queries 内の各 (i, x, limit) について、次を処理する
    • j <= m - 1 かつ nums[j] <= limit である間、以下を繰り返す
      • nums[j] を k ビット右シフトした値を prefixes に追加する
      • j := j + 1
    • prefixes が空の場合
      • res[i] := -1
    • それ以外の場合
      • res[i] を 1 ビット左シフトする(res[i] を 2 倍する)
      • target := res[i] XOR 1
      • x を k ビット右シフトした値と target の XOR が prefixes に存在する場合
        • res[i] := target
• res を返す

Pythonでの実装例

理解を深めるために、以下の実装を見てみましょう。

def solve(nums, queries):
    m, n = len(nums), len(queries)
    queries = sorted(((i, x, limit) for i, (x, limit) in enumerate(queries)), key=lambda x: x[2])
    nums = sorted(nums)
    res = [0] * n
    for k in range(31, -1, -1):
        prefixes = set()
        j = 0
        for i, x, limit in queries:
            while j <= m - 1 and nums[j] <= limit:
                prefixes.add(nums[j] >> k)
                j += 1
            if not prefixes:
                res[i] = -1
            else:
                res[i] <<= 1
                target = res[i] ^ 1
                if (x >> k) ^ target in prefixes:
                    res[i] = target
    return res

nums = [0,1,2,3,4]
queries = [[3,1],[1,3],[5,6]]
print(solve(nums, queries))

入力

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

出力

[3, 3, 7]

計算量の目安

クエリを limit 順に処理することで、ポインタ j は各ビット桁の処理中に後戻りせず前へ進むだけで済みます。全体の計算量は、ソートのコストに加えて約 32 ビット分の走査が必要になるため、おおよそ O((m + n) × 32 + m log m + n log n) 程度となります。これにより、大規模な入力でも高速に動作する効率的なアルゴリズムになっています。