Pythonでk番目に大きいXOR座標値を求めるプログラムの解説

問題の概要

m × n の行列と整数 k が与えられたとします。このとき、座標 (a, b) の値は、0 ≤ i ≤ a かつ 0 ≤ j ≤ b を満たすすべての要素 matrix[i][j] の XOR 値として定義されます。私たちの課題は、行列内の全座標の値の中からk 番目に大きい値(1始まりのインデックス)を見つけることです。

具体例

たとえば、次のような 2 × 2 の行列を考えてみましょう。

k = 1 の場合、答えは 7 になります。座標 (0, 1) の値は 5 XOR 2 = 7 と計算され、これが全座標の中で最大の値だからです。

解法のアプローチ

この問題は、「二次元累積 XOR」と「上位 k 個の値の管理」を組み合わせることで効率的に解けます。手順は以下の通りです。

1. 行数・列数の取得: m を行数、n を列数として設定します。
2. 行方向の累積 XOR: まず各行について、左から右へ隣接要素との XOR を順に計算します(matrix[i][j] = matrix[i][j] XOR matrix[i][j-1])。
3. 記録用辞書の準備: 値の出現回数を記録する辞書 seen を用意し、カウンター count を 0 で初期化します。
4. 列方向の累積 XOR: 次に各列について、上から下へ同様の累積 XOR を計算します。この二段階の処理により、各座標には原点 (0, 0) からその座標までの長方形領域全体の XOR 値が格納されます。
5. 上位 k 個のみ保持: 計算した値を seen に記録し、count が k を超えるたびに seen 内の最小値を取り除きます。これにより、常にそれまでの上位 k 個の値だけが残ります。
6. 結果の取得: すべての座標を処理し終えた後、seen に残っている最小値が求める「k 番目に大きい値」です。

Pythonでの実装例

理解を深めるために、実際の実装コードを見てみましょう。

def solve(matrix, k):
    m, n = len(matrix), len(matrix[0])

    # 行方向の累積XOR
    for i in range(m):
        for j in range(n):
            if j:
                matrix[i][j] ^= matrix[i][j-1]

    seen = {}
    count = 0

    # 列方向の累積XORと上位k個の管理
    for i in range(n):
        for j in range(m):
            if j:
                matrix[j][i] ^= matrix[j-1][i]

            seen[matrix[j][i]] = seen.get(matrix[j][i], 0) + 1
            count += 1

            if count > k:
                min_value = min(seen)
                seen[min_value] -= 1
                if not seen[min_value]:
                    seen.pop(min_value)

    return min(seen)

matrix = [[5,2],[1,6]]
k = 1
print(solve(matrix, k))

入力

[[5,2],[1,6]], 1

出力

7

処理のポイント

• 時間計算量: 二次元累積 XOR の計算に O(m×n)、count が k を超えてからの最小値探索 min(seen) に最大 O(k) ずつかかるため、全体の計算量は O(m×n×k) となります。heapq(優先度付きキュー)を利用すれば O(m×n×log k) まで高速化できます。
• 空間計算量: 常に上位 k 個分の値しか保持しないため、追加の記憶領域は O(k) で抑えられます。
• XOR の性質: XOR は同じ値を偶数回適用すると打ち消されるという性質を持つため、行方向・列方向の二段階の累積計算によって、各座標の長方形領域の XOR 値が正しく求まります。