Maximum total subarray value ii

impl Solution {
  pub fn max_total_value(nums: Vec<i32>, k: i32) -> i64 {
    // Optimized O(n * log(max_val)) algorithm:
    //
    // Binary search for v* = the k-th largest subarray value. For a given threshold v:
    //   count_le(v)  = # subarrays with (max - min) <= v  [O(n) sliding window]
    //   sum_le(v)    = sum of (max-min) for those subarrays [O(n) augmented deques]
    //
    // After finding v*:
    //   count_gt = total - count_le(v*)          [# subarrays strictly above v*]
    //   sum_gt   = sum_all - sum_le(v*)           [their total value]
    //   answer   = sum_gt + v* * (k - count_gt)
    //
    // count_le / sum_le use two monotonic deques (one for max, one for min).
    // Each deque stores (index, span-from-right); running sums updated in O(1) amortized.

    let n = nums.len() as i64;
    let k = k as i64;
    let total = n * (n + 1) / 2;

    if nums.is_empty() {
      return 0;
    }

    // --- helper: sliding-window count + sum of subarrays with max-min <= threshold ---
    let count_and_sum = |threshold: i64| -> (i64, i64) {
      let m = nums.len();
      // Deques store indices; max_deq is non-increasing in value, min_deq non-decreasing.
      // sum_max = sum over l in [left,right] of max([l,right])
      // sum_min = sum over l in [left,right] of min([l,right])
      let mut max_deq: std::collections::VecDeque<usize> = std::collections::VecDeque::new();
      let mut min_deq: std::collections::VecDeque<usize> = std::collections::VecDeque::new();
      let mut sum_max: i64 = 0;
      let mut sum_min: i64 = 0;
      let mut cnt: i64 = 0;
      let mut val_sum: i64 = 0;
      let mut left = 0usize;

      for right in 0..m {
        let rv = nums[right] as i64;

        // Extend max_deq: pop back while nums[back] <= rv, then push right.
        // Each popped entry back covered l in [prev+1, back]; remove its contribution.
        while !max_deq.is_empty() && nums[*max_deq.back().unwrap()] as i64 <= rv {
          let back = max_deq.pop_back().unwrap();
          let prev = max_deq.back().copied().map_or(left as i64 - 1, |x| x as i64);
          sum_max -= nums[back] as i64 * (back as i64 - prev);
        }
        let prev = max_deq.back().copied().map_or(left as i64 - 1, |x| x as i64);
        sum_max += rv * (right as i64 - prev);
        max_deq.push_back(right);

        // Extend min_deq: pop back while nums[back] >= rv, then push right.
        while !min_deq.is_empty() && nums[*min_deq.back().unwrap()] as i64 >= rv {
          let back = min_deq.pop_back().unwrap();
          let prev = min_deq.back().copied().map_or(left as i64 - 1, |x| x as i64);
          sum_min -= nums[back] as i64 * (back as i64 - prev);
        }
        let prev = min_deq.back().copied().map_or(left as i64 - 1, |x| x as i64);
        sum_min += rv * (right as i64 - prev);
        min_deq.push_back(right);

        // Shrink from left while window violates max - min > threshold.
        // Removing l=left: subtract its max and min contributions; pop front if stale.
        while nums[*max_deq.front().unwrap()] as i64
          - nums[*min_deq.front().unwrap()] as i64
          > threshold
        {
          sum_max -= nums[*max_deq.front().unwrap()] as i64;
          if *max_deq.front().unwrap() == left {
            max_deq.pop_front();
          }
          sum_min -= nums[*min_deq.front().unwrap()] as i64;
          if *min_deq.front().unwrap() == left {
            min_deq.pop_front();
          }
          left += 1;
        }

        let window = (right - left + 1) as i64;
        cnt += window;
        val_sum += sum_max - sum_min;
      }
      (cnt, val_sum)
    };

    // Binary search for smallest v* with count_le(v*) > total - k.
    let (_, sum_all) = {
      // compute sum_all via count_and_sum at max possible threshold (captures all subarrays)
      let max_val = *nums.iter().max().unwrap() as i64;
      let min_val = *nums.iter().min().unwrap() as i64;
      count_and_sum(max_val - min_val)
    };

    let need = total - k; // count_le(v*) must be > need → count_le(v*) >= need + 1

    let max_threshold = *nums.iter().max().unwrap() as i64 - *nums.iter().min().unwrap() as i64;
    let mut lo: i64 = 0;
    let mut hi: i64 = max_threshold;
    while lo < hi {
      let mid = lo + (hi - lo) / 2;
      let (c, _) = count_and_sum(mid);
      if c > need {
        hi = mid;
      } else {
        lo = mid + 1;
      }
    }
    let v_star = lo;

    let (count_le_vstar, sum_le_vstar) = count_and_sum(v_star);
    let count_gt = total - count_le_vstar;
    let sum_gt = sum_all - sum_le_vstar;

    sum_gt + v_star * (k - count_gt)
  }
}