Sum of imbalance numbers of all subarrays

Array Hash Table Enumeration

43.2% acceptance

Feb 25, 2026

325

The imbalance number of a 0-indexed integer array arr of length n is defined as the number of indices in sarr = sorted(arr) such that: 0 <= i < n - 1, and sarr[i+1] - sarr[i] > 1 Here, sorted(arr) is the function that returns the sorted version of arr. Given a 0-indexed integer array nums, return the sum of imbalance numbers of all its subarrays. A subarray is a contiguous non-empty sequence of elements within an array.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn sum_imbalance_numbers(nums: Vec<i32>) -> i32 {
    use std::collections::BTreeMap;
    let n = nums.len();
    let mut ans = 0i32;
    for l in 0..n {
      let mut freq: BTreeMap<i32, i32> = BTreeMap::new();
      let mut imb = 0i32;
      for r in l..n {
        let val = nums[r];
        let is_new = !freq.contains_key(&val);
        if is_new {
          // Find predecessor
          let pred = freq.range(..val).next_back().map(|(&k, _)| k);
          // Find successor (strictly greater)
          let succ = freq.range((val + 1)..).next().map(|(&k, _)| k);
          match (pred, succ) {
            (Some(p), Some(s)) => {
              let old = if s - p > 1 { 1 } else { 0 };
              let nl = if val - p > 1 { 1 } else { 0 };
              let nr = if s - val > 1 { 1 } else { 0 };
              imb += nl + nr - old;
            }
            (Some(p), None) => {
              imb += if val - p > 1 { 1 } else { 0 };
            }
            (None, Some(s)) => {
              imb += if s - val > 1 { 1 } else { 0 };
            }
            (None, None) => {}
          }
        }
        *freq.entry(val).or_insert(0) += 1;
        ans += imb;
      }
    }
    ans
  }
}