Maximum and minimum sums of at most size k subarrays

impl Solution {
  pub fn min_max_subarray_sum(nums: Vec<i32>, k: i32) -> i64 {
    let n = nums.len();
    let k = k as usize;
    // For each element as max: count subarrays of length 1..=k where it's the max
    // Use monotone stack approach
    // Similarly for min
    // Contribution of nums[i] as max:
    //   left[i] = distance to nearest element >= nums[i] to the left (exclusive)
    //   right[i] = distance to nearest element > nums[i] to the right (exclusive)
    //   (strict right to avoid double-counting)
    let mut left = vec![0usize; n]; // distance to left boundary or prev >= element
    let mut right = vec![0usize; n]; // distance to right boundary or next > element
    let mut stk: Vec<usize> = Vec::new();
    for i in 0..n {
      while stk.last().map_or(false, |&j| nums[j] < nums[i]) { stk.pop(); }
      left[i] = if let Some(&j) = stk.last() { i - j } else { i + 1 };
      stk.push(i);
    }
    stk.clear();
    for i in (0..n).rev() {
      while stk.last().map_or(false, |&j| nums[j] <= nums[i]) { stk.pop(); }
      right[i] = if let Some(&j) = stk.last() { j - i } else { n - i };
      stk.push(i);
    }
    let count_max = |i: usize| -> i64 {
      let l = left[i] as i64; let r = right[i] as i64;
      let ii = i as i64; let k = k as i64;
      let b_max = (ii + r - 1).min(ii + k - 1);
      // Split at threshold_b = ii + k - l:
      //   b <= threshold_b: a_min = ii+1-l, contribution = l each
      //   b >  threshold_b: a_min = b-k+1,  contribution = ii+k-b (arithmetic series)
      let threshold_b = ii + k - l;
      let n1 = (threshold_b.min(b_max) - ii + 1).max(0);
      let sum1 = l * n1;
      let b_start = (threshold_b + 1).max(ii);
      let sum2 = if b_max >= b_start {
        let n2 = b_max - b_start + 1;
        n2 * (ii + k) - (b_start + b_max) * n2 / 2
      } else { 0 };
      sum1 + sum2
    };
    // For min: use opposite convention (left = nearest >, right = nearest >=)
    let mut left2 = vec![0usize; n];
    let mut right2 = vec![0usize; n];
    stk.clear();
    for i in 0..n {
      while stk.last().map_or(false, |&j| nums[j] > nums[i]) { stk.pop(); }
      left2[i] = if let Some(&j) = stk.last() { i - j } else { i + 1 };
      stk.push(i);
    }
    stk.clear();
    for i in (0..n).rev() {
      while stk.last().map_or(false, |&j| nums[j] >= nums[i]) { stk.pop(); }
      right2[i] = if let Some(&j) = stk.last() { j - i } else { n - i };
      stk.push(i);
    }
    let count_min = |i: usize| -> i64 {
      let l = left2[i] as i64; let r = right2[i] as i64;
      let ii = i as i64; let k = k as i64;
      let b_max = (ii + r - 1).min(ii + k - 1);
      let threshold_b = ii + k - l;
      let n1 = (threshold_b.min(b_max) - ii + 1).max(0);
      let sum1 = l * n1;
      let b_start = (threshold_b + 1).max(ii);
      let sum2 = if b_max >= b_start {
        let n2 = b_max - b_start + 1;
        n2 * (ii + k) - (b_start + b_max) * n2 / 2
      } else { 0 };
      sum1 + sum2
    };
    let mut ans = 0i64;
    for i in 0..n {
      ans += nums[i] as i64 * count_max(i);
      ans += nums[i] as i64 * count_min(i);
    }
    ans
  }
}