Count prime gap balanced subarrays

impl Solution {
  pub fn prime_subarray(nums: Vec<i32>, k: i32) -> i32 {
    let max_val = 50001usize;
    // Sieve of Eratosthenes
    let mut is_prime = vec![true; max_val];
    is_prime[0] = false;
    is_prime[1] = false;
    let mut i = 2;
    while i * i < max_val {
      if is_prime[i] {
        let mut j = i * i;
        while j < max_val {
          is_prime[j] = false;
          j += i;
        }
      }
      i += 1;
    }

    let n = nums.len();
    // For each subarray, track min and max prime.
    // Use two pointers with monotone deques (like sliding window).
    // But we need: subarray has >=2 primes AND max_prime - min_prime <= k.
    // The constraint "max - min <= k" means we can use a sliding window if max-min is monotone.
    // However, max and min of primes in a subarray can increase or decrease.
    // Approach: For each right endpoint r, find the leftmost l_min such that
    // all subarrays [l..r] for l_min <= l <= r have max_prime - min_prime <= k and >= 2 primes.

    // Two-pointer: maintain window [left, right].
    // Count of valid subarrays ending at right = right - left + 1 - (subarrays with < 2 primes)
    // Actually use O(n^2) for 5*10^4? That might be 2.5*10^9 which is too slow.
    // Better approach: sliding window with deques for max and min prime values.
    // Maintain window where max_prime - min_prime <= k.
    // Additionally count must be >= 2.

    // Positions of primes in nums:
    let prime_positions: Vec<usize> = (0..n)
      .filter(|&i| is_prime[nums[i] as usize])
      .collect();

    if prime_positions.len() < 2 {
      return 0;
    }

    // For each right endpoint r (index in nums), let L(r) = leftmost index l such that
    // [l..=r] has max_prime - min_prime <= k.
    // Count subarrays [l..=r] that are valid = subarrays where l <= (second_prime_from_right)
    // meaning there are at least 2 primes in [l..=r].

    // We want: for each r, count l such that:
    // 1. max_prime[l..r] - min_prime[l..r] <= k
    // 2. At least 2 primes in [l..r]

    // Use two sliding windows:
    // window1: max_prime - min_prime <= k, gives [left1(r), r]
    // For condition 2: if the 2nd prime from right in prime_positions is at position p2,
    // then l <= p2 satisfies having >= 2 primes.

    // So valid l: left1(r) <= l <= min(r, p2_position)
    // where p2 = second last prime index <= r.

    let mut ans = 0i64;
    let mut left = 0usize;
    let mut max_dq: std::collections::VecDeque<usize> = std::collections::VecDeque::new();
    let mut min_dq: std::collections::VecDeque<usize> = std::collections::VecDeque::new();
    // prime_idx: index into prime_positions array for right boundary
    let _prime_ptr = 0usize; // next prime_positions index to process

    // For each r (in terms of nums indices), maintain the sliding window
    // We also need a pointer tracking the 2nd-from-last prime position.
    // Use a deque of prime positions within the window for that.

    let mut prime_in_window: std::collections::VecDeque<usize> = std::collections::VecDeque::new();

    for r in 0..n {
      // Update deques with nums[r] if it's prime
      if is_prime[nums[r] as usize] {
        // max_dq: decreasing queue of prime values
        while !max_dq.is_empty() && nums[*max_dq.back().unwrap()] <= nums[r] {
          max_dq.pop_back();
        }
        max_dq.push_back(r);
        // min_dq: increasing queue
        while !min_dq.is_empty() && nums[*min_dq.back().unwrap()] >= nums[r] {
          min_dq.pop_back();
        }
        min_dq.push_back(r);
        prime_in_window.push_back(r);
      }

      // Shrink window while max_prime - min_prime > k
      while !max_dq.is_empty() && !min_dq.is_empty()
        && nums[*max_dq.front().unwrap()] - nums[*min_dq.front().unwrap()] > k
      {
        left += 1;
        while max_dq.front().map_or(false, |&x| x < left) { max_dq.pop_front(); }
        while min_dq.front().map_or(false, |&x| x < left) { min_dq.pop_front(); }
        while prime_in_window.front().map_or(false, |&x| x < left) { prime_in_window.pop_front(); }
      }

      // Count valid l in [left, r] that have >= 2 primes in [l..r]
      // The 2nd prime from the right in the current window is prime_in_window[len-2] if exists.
      if prime_in_window.len() >= 2 {
        let second_last_prime = prime_in_window[prime_in_window.len() - 2];
        // Valid l: left <= l <= second_last_prime
        // count = second_last_prime - left + 1
        ans += second_last_prime as i64 - left as i64 + 1;
      }
    }

    ans as i32
  }
}