Competitive Programming Tutorial: ABC098D

Introduction

This article aims to improve both the author's and readers' algorithmic skills by carefully explaining some of the problems encountered during daily training.

Problem

Discussion

The following theorem holds:

For any non-negative integers $a$ and $b$, $a \oplus b \le a + b$ holds. Equality holds if and only if $a \land b = 0$.

From this theorem, when $a_l \oplus \cdots \oplus a_r = a_l + \cdots + a_r$ holds, the condition also holds for any sub-interval $[i, j]$ that satisfies $l \le i \le j \le r$. Conversely, when the interval $[l, r]$ does not satisfy the condition, $[l, i]$ will also not satisfy it for any $i > r$.

Therefore, the brute-force solution to this problem is as follows:

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
}

The code above will result in TLE as is, but with a slight improvement, the time complexity can be reduced to $O(AN)$.

When $a_i > 0$, the while loop terminates in at most $A = 20$ iterations. Therefore, it suffices to find the smallest $j > i$ such that $a_j > 0$ efficiently. Please refer to the solution example for the implementation.

Solution Example

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)
}