ABC091 D問題 : Two Sequeneces - 忘れても大丈夫

問題. Two Sequences

2つの長さ $n$ の非負数列 $a = a_1, \ldots, a_n$ ， $b = b_1, \ldots, b_n$ が与えられる． $\oplus_{1 \le i, j \le n} (a_i + b_j)$ を求めよ．ただし， $a_i \oplus b_j$ はビット単位の排他的論理和である．

制約： $1 \le n \le 200,000$ ，　 $0 \le a_i, b_i \le 2^{28}$

解法1. 二分探索

解説 [1] をそのまま参考にした．

ビット単位の排他的論理和はビット単位で独立なので，ビット単位でビットが立っているかどうかを調べる．0-indexed として， $k$ ビット目を調べる． $T = 2^k$ とすると，周期 $T$ で $k$ ビット目が立っているかどうかが分かる．つまり， $[0, T), [2T, 3T), [4T, 5T), \cdots$ のいずれかに含まれる自然数は $k$ ビット目が0で， $[T, 2T), [3T, 4T), [6T, 7T), \cdots$ のいずれかに含まれる自然数は $k$ ビット目が1である．したがって， $a, b$ の各要素をそれぞれ自身の値の $2T$ による剰余に置き換えると，どの組の和も $0 \le a_i + b_j \le 4T$ を満たす．次に， $b$ を昇順にソートする． $a_i$ を固定したときに， $a_i$ との和が $[T, 2T)$ または $[3T, 4T)$ に含まれる $b$ の要素数を数えると，そのパリティと $(a_i + b_1) \oplus (a_i + b_2) \oplus \cdots \oplus (a_i + b_n)$ の $k$ ビット目のパリティが一致する．よって，各 $a_i (1 \le i \le n)$ の $[T - a_i, 2T - a_i)$ または $[3T - a_i, 4T - a_i)$ に含まれる $b$ の要素数の総和のパリティが求める値の $k$ ビット目のパリティとなる．各ビット毎に $b$ のソートが支配的となるために計算時間は $O(n \log n)$ となる．また，答えとなるビット数の上限は $2 \times (2^{28} - 1) = 2^{29} - 2$ となり高々30ビットでint型に収まる．しかし，各 $k$ ビット目毎の区間に含まれる $b$ の要素数の総和の上限は $n^2$ となるため int型に収まらないので注意が必要である．

計算時間： $O(n \log n)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    int n;
    cin >> n;

    vector<int> a(n), b(n);
    for (auto &x : a) cin >> x;
    for (auto &x : b) cin >> x;

    const int size = log2(*max_element(a.begin(), a.end())
                          + *max_element(b.begin(), b.end()));

    int ans = 0;
    vector<int> bb(n);
    for (int i = 0, t = 1; i <= size; ++i) {
        for (int j = 0; j < n; ++j)
            bb[j] = b[j] % (2 * t);

        sort(bb.begin(), bb.end());

        ll cnt = 0;
        for (int j = 0; j < n; ++j) {
            const int aa = a[j] % (2 * t);

            // a_j + b_k \in [T, 2T)
            cnt += distance(lower_bound(bb.begin(), bb.end(), t - aa),
                            lower_bound(bb.begin(), bb.end(), 2 * t - aa));

            // a_j + b_k \in [3T, 4T)
            cnt += distance(lower_bound(bb.begin(), bb.end(), 3 * t - aa),
                            lower_bound(bb.begin(), bb.end(), 4 * t - aa));
        }

        // Update
        ans += (cnt % 2) * t;
        t *= 2;
    }

    cout << ans << endl;

    return 0;
}

解法2. ループ展開

[2] で参照されているように，ループ展開やSIMDを使うと愚直な二重ループの $O(n^2)$ 時間解法が通る．

次をソースコードに埋め込むとループ展開のオプション -funroll-loops と最適化オプション -O3 が追加される．

#pragma GCC optimize ("-O3", "unroll-loops")

全体でTLEとなるが， $n = 100,000$ のテストケースでTLEだったのが 1685 [ms] と通った．

ソースコードを表示 : ループ展開を使用した二重ループの解法(TLE)

#include <bits/stdc++.h>

using namespace std;

#pragma GCC optimize ("-O3", "unroll-loops")

constexpr int MAX_N = 200000;
int n, a[MAX_N], b[MAX_N];

int main() {
    // Input
    scanf("%d", &n);
    for (int i = 0; i < n; ++i) scanf("%d", &a[i]);
    for (int i = 0; i < n; ++i) scanf("%d", &b[i]);

    // Solve
    int ans = 0;
    for (int i = 0; i < n; ++i)
        for (int j = 0; j < n; ++j)
            ans ^= a[i] + b[j];

    // Output
    printf("%d\n", ans);

    return 0;
}

[2] のループ展開を参考に高速化を行った．ポイントとしては，ループ展開されやすいように書き直すのと， int型から unsigned型に変更することである．unsigned型への変更によって $n = 100,000$ のとき， 757 [ms] から 694 [ms] に少しだけの高速化となったが，これが TLE だった他のケースの AC に繋がった（ $n = 150,000$ では 1545 [ms] から 1701 [ms]）．下にACしたソースコードを載せる．ちなみに，bのブロックサイズを8192から4096に変更すると少し遅くなり，aのブロック化を無くすと TLE となった．
ループ展開の添字でブロックサイズが $2^k$ のとき，残った部分の開始時の添字は $n$ & ~ $(2^k - 1)$ となる．これは， $n$ から $n$ の $2^k$ による剰余を引いた値となるので， $n$ と同じブロックの先頭の要素と等しくなる（どこで参考にしたのか忘れてしまったが面白かったのでメモ）．

ソースコードを表示 : ループ展開を使用した二重ループの解法(AC)

#include <bits/stdc++.h>

using namespace std;

#pragma GCC optimize ("-O3", "unroll-loops")

constexpr int MAX_N = 200000;
size_t n;
unsigned a[MAX_N], b[MAX_N];

int main() {
    // Input
    scanf("%lu", &n);
    for (size_t i = 0; i < n; ++i) scanf("%u", &a[i]);
    for (size_t i = 0; i < n; ++i) scanf("%u", &b[i]);


    // Solve
    constexpr size_t block_a = 4, block_b = 8192;
    unsigned ans = 0;
    for (size_t i = 0; i + (block_a - 1) < n; i += block_a) {
        for (size_t j = 0; j + (block_b - 1) < n; j += block_b)
            for (size_t k = j; k < j + block_b; ++k) {
                ans ^= (a[i] + b[k]) ^ (a[i + 1] + b[k]);
                ans ^= (a[i + 2] + b[k]) ^ (a[i + 3] + b[k]);
            }

        for (size_t j = n & ~(block_b - 1); j < n; ++j) {
            ans ^= (a[i] + b[j]) ^ (a[i + 1] + b[j]);
            ans ^= (a[i + 2] + b[j]) ^ (a[i + 3] + b[j]);
        }
    }

    // rest of iterations
    for (size_t i = n & ~(block_a - 1); i < n; ++i) {
        for (size_t j = 0; j + (block_b - 1) < n; j += block_b)
            for (size_t k = j; k < j + block_b; ++k)
                ans ^= a[i] + b[k];

        for (size_t j = n & ~(block_b - 1); j < n; ++j)
            ans ^= a[i] + b[j];
    }

    // Output
    printf("%u\n", ans);

    return 0;
}