Maximum subarray xor with bounded range

impl Solution {
  pub fn max_xor(nums: Vec<i32>, k: i32) -> i32 {
    use std::collections::VecDeque;
    let n = nums.len();

    // Use prefix XOR + trie with sliding window.
    // prefix[0] = 0, prefix[i] = nums[0] ^ nums[1] ^ ... ^ nums[i-1]
    // XOR of subarray [l..=r] = prefix[r+1] ^ prefix[l]
    // Valid subarray [l..=r]: max(nums[l..=r]) - min(nums[l..=r]) <= k
    // For each r, let L(r) = smallest l such that [l..=r] is valid.
    // We want max over all r of: max over l in [L(r)..=r] of prefix[r+1] ^ prefix[l]
    //
    // We maintain a trie of prefix values. As r increases, L(r) can only increase.
    // We add prefix[r+1] and potentially remove old prefix values as L increases.

    let mut prefix = vec![0i32; n + 1];
    for i in 0..n {
      prefix[i + 1] = prefix[i] ^ nums[i];
    }

    // Trie node: children[0], children[1], count
    // We use 15 bits since nums[i] < 2^15, so prefix values < 2^15
    const BITS: usize = 15;
    let max_nodes = (n + 1) * (BITS + 1) + 2;
    let mut trie_left = vec![0usize; max_nodes];
    let mut trie_right = vec![0usize; max_nodes];
    let mut trie_count = vec![0i32; max_nodes];
    let mut trie_size = 1; // root is node 0

    let add = |val: i32, delta: i32,
           trie_left: &mut Vec<usize>, trie_right: &mut Vec<usize>,
           trie_count: &mut Vec<i32>, trie_size: &mut usize| {
      let mut node = 0;
      for bit in (0..BITS).rev() {
        let b = ((val >> bit) & 1) as usize;
        let next = if b == 0 {
          if trie_left[node] == 0 {
            let id = *trie_size;
            *trie_size += 1;
            trie_left[node] = id;
          }
          trie_left[node]
        } else {
          if trie_right[node] == 0 {
            let id = *trie_size;
            *trie_size += 1;
            trie_right[node] = id;
          }
          trie_right[node]
        };
        node = next;
        trie_count[node] += delta;
      }
    };

    let query_max = |val: i32,
             trie_left: &Vec<usize>, trie_right: &Vec<usize>,
             trie_count: &Vec<i32>| -> i32 {
      let mut node = 0;
      let mut result = 0;
      for bit in (0..BITS).rev() {
        let b = ((val >> bit) & 1) as usize;
        // We want the opposite bit to maximize XOR
        let want = 1 - b;
        let want_node = if want == 0 { trie_left[node] } else { trie_right[node] };
        if want_node != 0 && trie_count[want_node] > 0 {
          result |= 1 << bit;
          node = want_node;
        } else {
          let same_node = if b == 0 { trie_left[node] } else { trie_right[node] };
          node = same_node;
        }
      }
      result
    };

    // Sliding window for max - min <= k
    let mut max_deque: VecDeque<usize> = VecDeque::new();
    let mut min_deque: VecDeque<usize> = VecDeque::new();
    let mut left = 0usize;
    let mut ans = 0;

    // Add prefix[0] to trie (for subarray starting at index 0)
    add(prefix[0], 1, &mut trie_left, &mut trie_right, &mut trie_count, &mut trie_size);

    for right in 0..n {
      // Maintain deques for nums[left..=right]
      while !max_deque.is_empty() && nums[*max_deque.back().unwrap()] <= nums[right] {
        max_deque.pop_back();
      }
      max_deque.push_back(right);

      while !min_deque.is_empty() && nums[*min_deque.back().unwrap()] >= nums[right] {
        min_deque.pop_back();
      }
      min_deque.push_back(right);

      // Shrink from left while invalid
      while !max_deque.is_empty() && !min_deque.is_empty()
        && nums[*max_deque.front().unwrap()] - nums[*min_deque.front().unwrap()] > k
      {
        // Remove prefix[left] from trie
        add(prefix[left], -1, &mut trie_left, &mut trie_right, &mut trie_count, &mut trie_size);
        left += 1;
        while !max_deque.is_empty() && *max_deque.front().unwrap() < left {
          max_deque.pop_front();
        }
        while !min_deque.is_empty() && *min_deque.front().unwrap() < left {
          min_deque.pop_front();
        }
      }

      // Now valid subarrays ending at right have l in [left..=right]
      // XOR of [l..=right] = prefix[right+1] ^ prefix[l]
      // Trie contains prefix[left], prefix[left+1], ..., prefix[right]
      // Query max XOR with prefix[right+1]
      if left <= right + 1 {
        let val = query_max(prefix[right + 1], &trie_left, &trie_right, &trie_count);
        ans = ans.max(val);
      }

      // Add prefix[right+1] to trie for future use
      add(prefix[right + 1], 1, &mut trie_left, &mut trie_right, &mut trie_count, &mut trie_size);
    }

    ans
  }
}