yukicoder No.1240 - Or Sum of Xor Pair

No.1240 Or Sum of Xor Pair - yukicoder

問題

長さ $N$ の整数列 $(A_i)$ と整数 $X$ が与えられる。

$0 \leq i \lt j \lt N$ かつ $A_i \oplus A_j \lt X$ なる全ての $(i, j)$ について、 $A_i \lor A_j$ の総和を求めよ。

ここで $\lor, \oplus$ はそれぞれ bitwise or, bitwise xor を表す。

制約

$2 \leq N \leq 2 \cdot 10^5$
$1 \leq X \lt 2^{18} (=M)$
$0 \leq A_i \lt M$

考察

桁毎に考える、つまり

$0 \leq i \lt j \lt N$
$A_i \oplus A_j \lt X$
$A_i \lor A_j$ の $k$ 桁目が $1$

を満たす $(i, j)$ の個数を各 $k$ について求める。

ここで $A_i \lor A_j$ の $k$ 桁目が $1$ である $(i, j)$ の個数は、「全体」から「 $A_i, A_j$ の $k$ 桁目が共に $0$ であるものの個数」を引いたものとなる。よって $k$ 桁目が $0$ であるような $A_i$ のみを取り出せば、後者も前者と同じように数えられる。

後は

$0 \leq i \lt j \lt N$
$A_i \oplus A_j \lt X$

を満たす $(i, j)$ の個数を求めればよいが、これは俗に「xor 畳み込み」と呼ばれるテクニックによって $\Theta(N + M \log M)$ で数え上げられる。具体的なアルゴリズムはこちらを参照。

実装例

$i = j$ や $i \gt j$ なる $(i, j)$ の省き方に注意。 $i = j$ のとき $A_i \oplus A_j = 0$ であることを利用すると簡単に省ける。

#559055 (C++17) No.1240 Or Sum of Xor Pair - yukicoder

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81


#include <iostream>
#include <numeric>
#include <vector>

using lint = long long;
constexpr int K = 18;

// 自身とのxor畳み込み
void xor_convolution(std::vector<lint>& xs) {
    for (int i = 1; i < (1 << K); i <<= 1) {
        for (int b = 0; b < (1 << K); ++b) {
            if (b & i) continue;
            auto l = xs[b],
                 r = xs[b | i];
            xs[b] = l + r;
            xs[b | i] = l - r;
        }
    }

    for (auto& x : xs) x *= x;

    for (int i = 1; i < (1 << K); i <<= 1) {
        for (int b = 0; b < (1 << K); ++b) {
            if (b & i) continue;
            auto l = xs[b],
                 r = xs[b | i];
            xs[b] = l + r;
            xs[b | i] = l - r;
        }
    }

    for (auto& x : xs) x >>= K;
}

void solve() {
    int n, a;
    std::cin >> n >> a;

    std::vector<int> xs(n);
    for (auto& x : xs) std::cin >> x;

    lint tot;  // 全体の個数
    {
        int sz = 0;
        std::vector<lint> cnt(1 << K, 0);
        for (auto x : xs) {
            ++cnt[x];
            ++sz;
        }

        xor_convolution(cnt);
        cnt[0] -= sz;                 // 0にi=jなるものが含まれることに注意
        for (auto& c : cnt) c >>= 1;  // i>jを除外
        tot = std::accumulate(cnt.begin(), cnt.begin() + a, 0LL);
    }

    lint ans = 0;
    for (int k = 1; k < (1 << K); k <<= 1) {
        // lg(k)桁目が0であるもののみの個数
        lint num;

        {
            int sz = 0;
            std::vector<lint> cnt(1 << K, 0);
            for (auto x : xs) {
                if (x & k) continue;
                ++cnt[x];
                ++sz;
            }

            xor_convolution(cnt);
            cnt[0] -= sz;
            for (auto& c : cnt) c >>= 1;
            num = std::accumulate(cnt.begin(), cnt.begin() + a, 0LL);
        }

        ans += k * (tot - num);
    }

    std::cout << ans << "\n";
}