ABC117D - XXOR

問題概要

リンク

$N$個の非負整数$A_1,A_2,\cdots A_N$と$K$が与えられる。

$f(X) = X \oplus A_1 + X \oplus A_2 + \cdots + X \oplus A_N$とするとき、$0 \leq X \leq K$での$f$の最大値を求めよ。

制約

$0 \leq N \leq 10^5$

$0 \leq K \leq 10^{12}$

$0 \leq A_i \leq 10^{12}$

解法

$f(X)$を求めるために$0 \leq X \leq K$の範囲を全て計算して最も大きくなるものを求めたいが、制約上$O(K)$では TLE になってしまう。そこで、$X$ に入る値を全て調べることなく解にたどり着けるようなアプローチを考える。

$X \leq K$を無視して考えてみる

まず、$X \leq K$の制限がない場合、$f(X)$が最大になるには$X$がどのようになっていればいいか考える。$f(X)$の性質上、$X$を 2 進数で表した時の$i$桁目の値は各$A_1,...A_N$の$i$桁目の値にしか関係がないので、各桁ごとに独立に考えることができる。

$X$を 2 進数で表した時の$i$桁目($X_i$と表記する)が$f(X)$に寄与する値$f_{X_i}$に注目すると、

$f_{X_i} = \Biggl\{ \begin{array}{cc} X_iが0のとき & A_1,...A_Nのうち、 i桁目が1になっているものの数 \times 2^{i-1} \\ X_iが1のとき & A_1,...A_Nのうち、 i桁目が0になっているものの数 \times 2^{i-1} \\ \end{array}$$$

となる。$X$が 2 進数で 60 桁まで取れるとすると、$f(X) = \sum_{i=1}^{60} f_{X_i}$なので、

あらかじめ各$i$について$A_1,...A_N$の$i$桁目が 1 になっている個数をカウントしておけば、X の$i$桁目は$f_{X_i}$がより大きいものを選んでいけばいいので$O(N)$で解くことができる。

$X \leq K$を考慮する

$f(X)$を考えるにあたり、$X$ の各ビット毎に独立で考えることができることが分かったが、今回は厄介な$X \leq K$という制約がある。この制約により、ある$i$に対する$X_i$を自由に決めてよいかどうかが、$i$以外の他のビットをどのように決めたかに依存してしまう。

そこで、他のビットがどうなっている状態なら$X_i$ を自由に決めてよいかを考察してみる。

まず、$K$の$i$桁目($K_i$)が 0 なら、そのビットを 1 にした$K_{+i}$を作った場合、明らかに$K_{+i} > K$なので$X=K_{+}$とできない。 一方$K_i = 1$なら、そのビットを 0 にした$K_{-i}$は$K > K_{-i}$になるため、$X=K_{-i}$とすることができる。

また、$K_i = 1$のときに$X=K_{-i}$としたとき、その後$i$より下位の桁を全て 1 に変更したとしても$X \leq K$が保たれる。

まとめると、以下のようになる。

$X_i = \Biggl\{ \begin{array}{cc} K_i = 0 のとき & iよりも上位のビットjで、K_j=1かつX_j=0なら 0か1、そうでないなら0 \\ K_i = 1 のとき & 0か1 \\ \end{array}$$$

このような制約の元で$X_i$を決めるには、$X=K$としたときと、$K_i=1$となるビットで$X_i=0$として、下位ビット全ておよび上位の$K_j=1$となっているビットを自由に選んだものを検証すればよい。

例として、$K=11001101$ としたとき、自由に選んでよいビットを?で表すと

$\begin{array}{cc} 11001101 & (X=K) \\ ??00??00 & (K_1=0) \\ ??00?0?? & (K_3=0) \\ ??000??? & (K_4=0) \\ ?0?????? & (K_7=0) \\ 0??????? & (K_8=0) \\ \end{array}$$$

を調べて、$f(X)$が最大になる$X_i$を選べばよい。

Code

下記が、rust でこの問題を解いた場合のコードです。

fn solve() {
    input! {
        n: usize,
        k: usize,
        a: [usize;n]
    }

    let mut k_bit = vec![0; 60];
    for b in 0..60 {
        k_bit[b] = k >> b & 1;
    }

    let mut a_sum_bit = vec![0; 60];
    for i in 0..n {
        let a_i = a[i];
        for b in 0..60 {
            a_sum_bit[b] += a_i >> b & 1;
        }
    }

    // まずx=kの場合のf(x)を見る
    let mut ans = 0;
    for i in 0..60 {
        let base = 1 << i;
        ans += if k_bit[i] == 1 {
            (a.len() - a_sum_bit[i]) * base
        } else {
            a_sum_bit[i] * base
        }
    }

    for i in 0..60 {
        let mut tmp_ans = 0;

        // kのi+1ビット目が1なら、そのビットを0にする代わりに下位のビットは全て0か1のより大きい方として良い
        if k_bit[i] == 1 {
            for j in 0..60 {
                let base = 1 << j;
                let bit_j_0 = a_sum_bit[j] * base;
                let bit_j_1 = (a.len() - a_sum_bit[j]) * base;
                if j < i {
                    tmp_ans += cmp::max(bit_j_0, bit_j_1);
                } else if j == i {
                    tmp_ans += bit_j_0;
                } else {
                    tmp_ans += if k_bit[j] == 0 { bit_j_0 } else { bit_j_1 };
                }
            }
        }

        ans = cmp::max(tmp_ans, ans);
    }

    println!("{}", ans);
}