LeetCode 3632 - Subarrays with XOR at Least K

This problem asks us to count the number of contiguous subarrays of a given array nums such that the bitwise XOR of all elements in the subarray is greater than or equal to a given integer k. The input array nums consists of positive integers, and k is a non-negative integer.

leetcodehardarraybit-manipulationtrieprefix-sum

LeetCode Problem 3632

Difficulty: 🔴 Hard
Topics: Array, Bit Manipulation, Trie, Prefix Sum

Solution

Problem Understanding

This problem asks us to count the number of contiguous subarrays of a given array nums such that the bitwise XOR of all elements in the subarray is greater than or equal to a given integer k. The input array nums consists of positive integers, and k is a non-negative integer. A contiguous subarray means the elements are consecutive in the original array, and XOR refers to the standard bitwise exclusive OR operation.

For instance, in the example nums = [3,1,2,3], k = 2, we need to consider all possible subarrays like [3], [3,1], [1,2], [2,3], etc., calculate their XOR, and count how many meet or exceed the threshold k.

The constraints are significant: nums can have up to 10^5 elements, and each element can be as large as 10^9. This immediately rules out any solution that examines all O(n²) subarrays directly, as it would be too slow. Edge cases include arrays of size 1, arrays of zeros, and k = 0, where every subarray satisfies the condition.

Approaches

Brute Force

A naive approach is to iterate through all possible subarrays. For each starting index i, iterate through every ending index j ≥ i, calculate the XOR of elements between i and j, and check if it is ≥ k. While correct, this approach requires calculating XOR for O(n²) subarrays, and each XOR calculation may take O(n) in the worst case if computed naively (though it can be optimized using prefix XORs). Even with prefix XORs, the brute force would still be O(n²), which is too slow for n up to 10^5.

Optimal Approach

The key insight is to use prefix XORs combined with a Trie (binary prefix tree) to efficiently count subarrays that satisfy the XOR condition. Define the prefix XOR array prefixXor[i] as the XOR of the first i elements. The XOR of a subarray [l..r] can be computed as prefixXor[r] ^ prefixXor[l-1]. The problem then reduces to counting, for each prefixXor[r], the number of prefixXor[l-1] values such that (prefixXor[r] ^ prefixXor[l-1]) >= k. A binary trie allows counting this efficiently in O(30) per query since integers are at most 30 bits.

Approach Time Complexity Space Complexity Notes
Brute Force O(n²) O(n) Iterate all subarrays using prefix XOR, check each
Optimal O(n * 30) O(n * 30) Use prefix XOR + Trie to count XORs efficiently

Algorithm Walkthrough

  1. Initialize a Trie that will store binary representations of prefix XORs. Each node in the Trie tracks two children (0 and 1) and a count of how many numbers pass through that node.
  2. Initialize prefixXor = 0 and insert 0 into the Trie to handle subarrays starting at index 0.
  3. Iterate through each number in nums. Update the prefixXor as prefixXor ^= num.
  4. For the current prefixXor, query the Trie to count how many previous prefix XORs p satisfy (prefixXor ^ p) >= k. This is done by traversing the Trie from the most significant bit, deciding whether the bit contributes to a valid XOR count at each position.
  5. Add the result of the query to a running total ans.
  6. Insert the current prefixXor into the Trie to consider it for future subarrays.
  7. Return ans after processing all elements.

Why it works: The prefix XOR trick allows us to transform the subarray XOR problem into a problem of counting previous XORs with a specific property. The Trie efficiently tracks the prefix XORs in binary form, allowing O(30) traversal per element, making the overall approach linear in practice for up to 10^5 elements.

Python Solution

from typing import List

class TrieNode:
    def __init__(self):
        self.children = [None, None]
        self.count = 0

class Trie:
    def __init__(self):
        self.root = TrieNode()
    
    def insert(self, num: int) -> None:
        node = self.root
        for i in reversed(range(30)):
            bit = (num >> i) & 1
            if not node.children[bit]:
                node.children[bit] = TrieNode()
            node = node.children[bit]
            node.count += 1

    def query(self, num: int, k: int) -> int:
        node = self.root
        res = 0
        for i in reversed(range(30)):
            if not node:
                break
            num_bit = (num >> i) & 1
            k_bit = (k >> i) & 1
            if k_bit == 1:
                if node.children[num_bit]:
                    res += node.children[num_bit].count
                node = node.children[1 - num_bit]
            else:
                node = node.children[num_bit]
        return res

class Solution:
    def countXorSubarrays(self, nums: List[int], k: int) -> int:
        trie = Trie()
        trie.insert(0)
        prefixXor = 0
        ans = 0
        for num in nums:
            prefixXor ^= num
            ans += trie.query(prefixXor, k)
            trie.insert(prefixXor)
        return ans

Implementation Walkthrough: We define a TrieNode class to store children and count, and a Trie class for insertion and query. Each prefix XOR is inserted into the Trie. The query function traverses the Trie based on bits of the current prefix XOR and k to efficiently count how many previous prefixes produce a subarray XOR ≥ k. Finally, we return the accumulated answer.

Go Solution

package main

func countXorSubarrays(nums []int, k int) int64 {
	type TrieNode struct {
		children [2]*TrieNode
		count    int
	}
	root := &TrieNode{}
	var insert func(num int)
	insert = func(num int) {
		node := root
		for i := 29; i >= 0; i-- {
			bit := (num >> i) & 1
			if node.children[bit] == nil {
				node.children[bit] = &TrieNode{}
			}
			node = node.children[bit]
			node.count++
		}
	}

	var query func(num, k int) int
	query = func(num, k int) int {
		node := root
		res := 0
		for i := 29; i >= 0; i-- {
			if node == nil {
				break
			}
			numBit := (num >> i) & 1
			kBit := (k >> i) & 1
			if kBit == 1 {
				if node.children[numBit] != nil {
					res += node.children[numBit].count
				}
				node = node.children[1^numBit]
			} else {
				node = node.children[numBit]
			}
		}
		return res
	}

	insert(0)
	prefixXor := 0
	var ans int64
	for _, num := range nums {
		prefixXor ^= num
		ans += int64(query(prefixXor, k))
		insert(prefixXor)
	}
	return ans
}

Go Notes: Go requires explicit handling of pointers for Trie nodes. The XOR computation and insertion logic is identical, but we ensure int64 accumulation to avoid overflow. The ^ operator is used for XOR in Go.

Worked Examples

Example 1: nums = [3,1,2,3], k = 2

Step prefixXor Trie Nodes (binary inserted) Subarrays Counted ans
0 0 0 - 0
1 3 0,3 [3] 1
2 2 0,3,2 [3,1],[1,2] 3
3 0 0,3,2,0 [3,1,2,3] 4
4 3 ... [2],[3] 6

Example 2: nums = [0,0,0], k = 0

Every subarray has XOR = 0, so all 6 subarrays counted. Prefix XOR = [0,0,0,0].

Complexity Analysis

Measure Complexity Explanation
Time O(n * 30) Each number inserts and queries in Trie with 30 bits (max for 10^9)
Space O(n * 30) Trie stores up to n prefix XORs, each up to 30 bits

Even though Trie nodes scale with n, the constant factor is manageable due to the 30-bit limit.

Test Cases

# provided examples
assert Solution().countX