C++でビット単位ANDがゼロになるトリプルを数える方法

問題の概要

整数型の配列 A が与えられているとします。このとき、次の条件をすべて満たすインデックスのトリプル (i, j, k) の個数を求める必要があります。

• 0 <= i < A のサイズ
• 0 <= j < A のサイズ
• 0 <= k < A のサイズ

そして、A[i] AND A[j] AND A[k] の計算結果が 0 になることです。ここでいう AND は、ビットごとの論理積（bitwise-AND）演算子を表します。

たとえば、入力が [3, 1, 2] の場合、条件を満たすトリプルは合計 12 個存在するため、出力は 12 となります。

解法のアプローチ

すべてのトリプルを総当たりで調べると計算量が O(n³) となり、配列が大きくなるほど非効率です。そこで、あらかじめ「2つの要素のANDの結果」をハッシュマップに記録しておくことで、計算量を O(n²) まで削減できます。具体的な手順は以下の通りです。

1. マップ m を用意し、答えとなる変数 ret を 0 で初期化します。
2. n := A のサイズとします。
3. 二重ループですべてのペア (i, j) を走査し、A[i] & A[j] の値が出現する回数をマップ m にカウントします。
4. 次に各 i について x := A[i] とし、マップ内のすべてのキーと値のペア a を確認します。
5. (a.key AND x) が 0 に等しい場合、ret := ret + a.value を実行します。
6. 最後に ret を返します。

この方法では、「i と j を選んだときのAND結果」を再利用できるため、3つ目の要素との組み合わせを高速に判定できます。

C++による実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   int countTriplets(vector<int>& A){
      unordered_map<int, int> m;
      int ret = 0;
      int n = A.size();
      for (int i = 0; i < n; i++) {
         for (int j = 0; j < n; j++) {
            m[A[i] & A[j]]++;
         }
      }
      for (int i = 0; i < n; i++) {
         int x = A[i];
         for (auto& a : m) {
            if ((a.first & x) == 0) {
               ret += a.second;
            }
         }
      }
      return ret;
  }
};
main(){
   Solution ob;
   vector<int> v = {3,1,2};
   cout << (ob.countTriplets(v));
}

入力

{3,1,2}

出力

12

計算量について

最初のフェーズで全ペアのAND結果を集計するのに O(n²)、次のフェーズでマップのサイズを M とすると O(n × M) かかります。マップのエントリ数は取り得るAND値の種類に依存するため、要素の値が小さい場合には非常に効率的に動作します。素朴な O(n³) の総当たり法と比べると、大きな配列でも実用的な速度で処理できる点がこの手法の大きな利点です。