1442. Count Triplets That Can Form Two Arrays of Equal XOR solved in Rust

Link

1442. Count Triplets That Can Form Two Arrays of Equal XOR

Problem

The problem is asking to count the number of triplets (i, j, k) such that i < j <= k and arr[i] ^ arr[i+1] ^ ... ^ arr[j-1] == arr[j] ^ arr[j+1] ^ ... ^ arr[k].

Which means the XOR of the subarray arr[i..j] is equal to the XOR of the subarray arr[j..k+1].

Idea

To better understand the problem let's try to solve it using a brute force approach. We will treat XOR as any other math operation and try to find the triplets that satisfy the condition.

The brute force approach is to iterate over all the possible triplets and check if the XOR of the subarrays is equal. If it is, we will increment the result.

Time complexity of this approach is O(n^4). Highly inefficient.

// src/lib.rs

pub fn count_triplets(arr: Vec<i32>) -> i32 {
    let len = arr.len();
    let mut res = 0;

    for i in 0..len {
        for j in i + 1..len {
            for k in j..len {
                let mut a = 0;
                let mut b = 0;

                for idx in i..j {
                    a ^= arr[idx];
                }

                for idx in j..=k {
                    b ^= arr[idx]
                }

                if a == b {
                    res += 1
                }
            }
        }
    }

    res
}

Solution

In order to solve this problem efficiently, we need to understand the properties of the XOR operation. The XOR operation is commutative and associative. Which means a ^ b == b ^ a and a ^ (b ^ c) == (a ^ b) ^ c.

The XOR of a number with itself is 0. Which means a ^ a == 0.

// src/lib.rs

pub struct Solution;

impl Solution {
    pub fn count_triplets(arr: Vec<i32>) -> i32 {
        let len = arr.len();
        let mut res = 0;

        for i in 0..len {
            let mut cur_xor = arr[i];

            for k in i + 1..len {
                cur_xor ^= arr[k];

                if cur_xor == 0 {
                    res += (k - i) as i32
                }
            }
        }

        res
    }  
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn it_works() {
        let scenarios = vec![(vec![2, 3, 1, 6, 7], 4), (vec![1, 1, 1, 1, 1], 10)];

        scenarios
            .into_iter()
            .enumerate()
            .for_each(|(idx, (input, expected))| {
                let result = Solution::count_triplets(input);
                assert_eq!(result, expected);
                println!("  ✓ scenario {}", idx + 1)
            });
    }
}

Analysis

The time complexity of this solution is O(n^2) and the space complexity is O(1).