競プロ解説 [ABC098D]

はじめに

本記事は、日々の精進で出会った問題の一部を丁寧目に解説することで、筆者と読者のアルゴリズム力を向上させることを目的としています。

問題

考察

次の定理が成り立ちます。

任意の非負整数 $a, b$ について、$a \oplus b \le a + b$ が成り立つ。等号は $a \land b = 0$ のときに限り成り立つ。

定理より、$a_l \oplus\cdots\oplus a_r = a_l + \cdots a_r$ が成り立つとき、$l \le i \le j \le r$ を満足する任意の区間 $[i, j]$ について条件が成り立ちます。逆に、区間 $[l, r]$ が条件を満足しないとき、任意の $i > r$ について $[l, i]$ も条件を満足しません。

したがって、本問題の愚直な解法は次の通りです。

let mut counter = 0

for l in 0..n {
    let mut r = l + 1;
    let mut acc_sum = a[l];
    while r < n && acc_sum & a[r] == 0 {
        acc_sum += a[r];
        r += 1
    }

    counter += r - l
}

上記のコードはこのままでは TLE してしまいますが、少しの工夫で計算量を $O(AN)$ に抑えることができます。

$a_i > 0$ のとき、while ループは高々 $A = 20$ 回で終了します。したがって、$a_j > 0$ を満足する最小の $j > i$ を高速に見つけることができればよいです。実装は回答例を参照してください。

回答例

use itertools::Itertools;
use proconio::input;

fn main() {
    input! { n: usize, mut a: [u32; n], }

    let next_nonzero = {
        let mut res = (0..n)
            .rev()
            .scan(n, |state, i| {
                let res = Some(*state);
                if a[i] > 0 {
                    *state = i
                }
                res
            })
            .collect_vec();
        res.reverse();

        res
    };

    let mut res = 0;
    for l in 0..n {
        let mut r = next_nonzero[l];
        let mut acc_sum = a[l];
        // up to 20 steps
        while r < n && acc_sum & a[r] == 0 {
            acc_sum += a[r];
            r = next_nonzero[r]
        }

        res += r - l
    }

    println!("{}", res)
}