C++でXORが等しくなる2つの部分配列を持つトリプレットの個数を数える方法

問題概要

整数型の配列 arr が与えられたとします。ここで、(0 <= i < j <= k < N) を満たす3つのインデックス i、j、k を選ぶことを考えます(N は配列のサイズです)。このとき、a と b の値はそれぞれ次のように定義されます。

a = arr[i] XOR arr[i + 1] XOR ... XOR arr[j - 1]
b = arr[j] XOR arr[j + 1] XOR ... XOR arr[k]

この問題の目的は、a と b が等しくなるようなトリプレット (i, j, k) の総数を求めることです。

例として、入力が [2,3,1,6,7] の場合を考えてみましょう。このとき出力は 4 となります。条件を満たすトリプレットは (0,1,2)、(0,2,2)、(2,3,4)、(2,4,4) の4つです。

解法のアプローチ

この問題は、マップ(連想配列)を活用することで効率的に解くことができます。手順は以下の通りです。

• 結果を格納する変数 ret を 0 で初期化する

• n を配列 arr のサイズとする

• i を 1 から開始し、i < n の間、i を1ずつ増やしながら以下を繰り返す

  • マップ m を定義する

  • x1 := 0、x2 := 0 で初期化する

  • j を i - 1 から開始し、j >= 0 の間、j を1ずつ減らしながら、x1 に arr[j] の XOR を累積し、m[x1] のカウントを1増やす

  • j を i から開始し、j < n の間、j を1ずつ増やしながら、x2 に arr[j] の XOR を累積し、ret に m[x2] の値を加算する

• 最後に ret を返す

このアプローチのポイントは、境界位置 i ごとに、左側の部分配列で作られる各 XOR 値の出現回数をマップに記録しておき、右側の部分配列の XOR 値と照合することで、a = b となる組み合わせを効率よく数えられる点です。計算量は O(n²) となり、すべての組み合わせを素朴に調べる O(n³) の方法よりも高速に動作します。

実装例

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

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int countTriplets(vector<int>& arr) {
        int ret = 0;
        int n = arr.size();
        for (int i = 1; i < n; i++) {
            map<int, int> m;
            int x1 = 0;
            int x2 = 0;
            for (int j = i - 1; j >= 0; j--) {
                x1 = x1 ^ arr[j];
                m[x1]++;
            }
            for (int j = i; j < n; j++) {
                x2 = x2 ^ arr[j];
                ret += m[x2];
            }
        }
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {2,3,1,6,7};
    cout << (ob.countTriplets(v));
}

入力

{2,3,1,6,7}

出力

4