LeetCode 1442 - Count Triplets That Can Form Two Arrays of Equal XOR

LeetCode Problem 1442

Difficulty: 🟡 Medium
Topics: Array, Hash Table, Math, Bit Manipulation, Prefix Sum

Solution

Problem Understanding

The problem asks us to find the number of triplets (i, j, k) in an array arr such that the XOR of elements from index i to j-1 equals the XOR of elements from index j to k. Formally, for 0 <= i < j <= k < arr.length, define:

a = arr[i] ^ arr[i+1] ^ ... ^ arr[j-1]
b = arr[j] ^ arr[j+1] ^ ... ^ arr[k]

We need to count all triplets where a == b.

The input is a 1-dimensional integer array with length up to 300, and each element is at most 10^8. The output is a single integer representing the number of valid triplets.

Key observations include:

1. The XOR operation is associative and commutative, which allows us to use prefix XOR sums to compute ranges efficiently.
2. Edge cases include arrays with identical elements, very short arrays, or arrays where no triplet exists.
3. The constraints allow for an O(n^2) solution, but an O(n^3) brute-force approach may be too slow.

Approaches

Brute Force

The brute-force solution iterates over all possible triplets (i, j, k) and computes a and b directly by iterating over the subarrays. While correct, this approach has three nested loops, giving a time complexity of O(n^3). For n = 300, this could involve up to 27 million iterations, which is inefficient.

Optimal Approach

The key insight is using prefix XOR. Define prefixXOR[x] as the XOR of all elements from index 0 to x-1. Then, the XOR of any subarray arr[i:j] can be computed in O(1) as:

arr[i] ^ ... ^ arr[j] = prefixXOR[j+1] ^ prefixXOR[i]

We can rewrite the condition a == b as:

prefixXOR[j] ^ prefixXOR[i] == prefixXOR[k+1] ^ prefixXOR[j]

Simplifying gives:

prefixXOR[i] == prefixXOR[k+1]

This allows us to fix i and k and then compute the number of valid j values in between. Every pair (i, k) where prefixXOR[i] == prefixXOR[k+1] contributes (k - i) valid triplets. This reduces the time complexity to O(n^2).

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n^3)	O(1)	Check all triplets explicitly
Optimal	O(n^2)	O(n)	Use prefix XOR and count valid ranges

Algorithm Walkthrough

1. Initialize a prefixXOR array of length n+1, with prefixXOR[0] = 0.
2. Compute prefixXOR[i+1] = prefixXOR[i] ^ arr[i] for all i from 0 to n-1. This stores XOR of all elements up to but not including index i+1.
3. Initialize a counter count = 0.
4. Iterate over all i from 0 to n-1.
5. For each i, iterate over all k from i+1 to n-1.
6. If prefixXOR[i] == prefixXOR[k+1], add (k - i) to count.
7. Return count.

Why it works: XOR has the property that x ^ x = 0 and x ^ 0 = x. By computing prefix XOR, we reduce the problem of comparing two subarrays to a simple equality check between prefixXOR[i] and prefixXOR[k+1]. The number of valid j indices between i and k is exactly (k - i).

Python Solution

from typing import List

class Solution:
    def countTriplets(self, arr: List[int]) -> int:
        n = len(arr)
        prefixXOR = [0] * (n + 1)
        
        for i in range(n):
            prefixXOR[i + 1] = prefixXOR[i] ^ arr[i]
        
        count = 0
        for i in range(n):
            for k in range(i + 1, n):
                if prefixXOR[i] == prefixXOR[k + 1]:
                    count += (k - i)
        
        return count

Explanation: We first compute prefixXOR to allow O(1) subarray XOR computation. Then we iterate over all possible pairs (i, k) and check if they satisfy the condition. Each valid pair contributes (k - i) triplets.

Go Solution

func countTriplets(arr []int) int {
    n := len(arr)
    prefixXOR := make([]int, n+1)
    
    for i := 0; i < n; i++ {
        prefixXOR[i+1] = prefixXOR[i] ^ arr[i]
    }
    
    count := 0
    for i := 0; i < n; i++ {
        for k := i + 1; k < n; k++ {
            if prefixXOR[i] == prefixXOR[k+1] {
                count += k - i
            }
        }
    }
    
    return count
}

Go-specific notes: The slice prefixXOR is preallocated with make and initialized to zeros automatically. Integer overflow is not a concern as Go uses 64-bit integers by default for XOR operations on slices of int.

Worked Examples

Example 1: arr = [2,3,1,6,7]

Compute prefixXOR: [0, 2, 1, 0, 6, 1].

Iterate (i, k) pairs:

i	k	prefixXOR[i]	prefixXOR[k+1]	k-i	Add to count?
0	1	0	1	1	No
0	2	0	0	2	Yes, count+=2
0	3	0	6	3	No
0	4	0	1	4	No
1	2	2	0	1	No
1	3	2	6	2	No
1	4	2	1	3	No
2	3	1	6	1	No
2	4	1	1	2	Yes, count+=2
3	4	0	1	1	No

Total count = 4.

Example 2: arr = [1,1,1,1,1]

Compute prefixXOR: [0,1,0,1,0,1].

Following the same logic, every pair (i,k) where prefixXOR[i] == prefixXOR[k+1] contributes k-i triplets. Total count = 10.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n^2)	Two nested loops over n elements, computing subarray XOR in O(1) using prefixXOR
Space	O(n)	Store prefixXOR array of size n+1

The complexity is efficient for n ≤ 300, as O(n^2) involves at most 90,000 iterations, which is acceptable.

Test Cases

# Example cases
assert Solution().countTriplets([2,3,1,6,7]) == 4  # given example 1
assert Solution().countTriplets([1,1,1,1,1]) == 10  # given example 2

# Edge cases
assert Solution().countTriplets([1]) == 0  # single element, no triplets
assert Solution().countTriplets([1,2]) == 0  # two elements, no triplets
assert Solution().countTriplets([1,2,3]) == 1  # minimal triplet
assert Solution().countTriplets([0,0,0,0]) == 6  # all zeros, multiple triplets
assert Solution().countTriplets([1,2,1,2,1]) == 4  # alternating pattern

Test	Why
[2,3,1,6,7]	Validates general case with mix of numbers
[1,1,1,1,1]	All elements identical
[1]	Single element, no triplets possible
[1,2]	Array too short for a triplet
[1,2,3]	Minimal array forming exactly