Sum of Imbalance Numbers of All Subarrays

Problem

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 thesum of imbalance numbers of all its subarrays.

A subarray is a contiguous non-empty sequence of elements within an array.

Examples

Example 1

Input: nums = [2,3,1,4]
Output: 3
Explanation: There are 3 subarrays with non-zero**** imbalance numbers:
- Subarray [3, 1] with an imbalance number of 1.
- Subarray [3, 1, 4] with an imbalance number of 1.
- Subarray [1, 4] with an imbalance number of 1.
The imbalance number of all other subarrays is 0. Hence, the sum of imbalance numbers of all the subarrays of nums is 3. 

Example 2

Input: nums = [1,3,3,3,5]
Output: 8
Explanation: There are 7 subarrays with non-zero imbalance numbers:
- Subarray [1, 3] with an imbalance number of 1.
- Subarray [1, 3, 3] with an imbalance number of 1.
- Subarray [1, 3, 3, 3] with an imbalance number of 1.
- Subarray [1, 3, 3, 3, 5] with an imbalance number of 2. 
- Subarray [3, 3, 3, 5] with an imbalance number of 1. 
- Subarray [3, 3, 5] with an imbalance number of 1.
- Subarray [3, 5] with an imbalance number of 1.
The imbalance number of all other subarrays is 0. Hence, the sum of imbalance numbers of all the subarrays of nums is 8. 

Constraints

• 1 <= nums.length <= 1000
• 1 <= nums[i] <= nums.length

Solution

Approach

For each subarray, keep a set of seen numbers and track the imbalance as you extend the subarray. For each new number, update the imbalance by checking its neighbors. Since nums[i] <= n, use a boolean array for the set.

---

Code

C++

#include <vector>
using namespace std;
class Solution {
public:
    int sumImbalanceNumbers(vector<int>& nums) {
        int n = nums.size(), res = 0;
        for (int i = 0; i < n; ++i) {
            vector<bool> seen(n+2);
            seen[nums[i]] = true;
            int imb = 0;
            for (int j = i+1; j < n; ++j) {
                int x = nums[j];
                if (!seen[x]) {
                    if (!seen[x-1] && !seen[x+1]) imb++;
                    if (seen[x-1] && seen[x+1]) imb--;
                }
                seen[x] = true;
                res += imb;
            }
        }
        return res;
    }
};

Java

class Solution {
    public int sumImbalanceNumbers(int[] nums) {
        int n = nums.length, res = 0;
        for (int i = 0; i < n; ++i) {
            boolean[] seen = new boolean[n+2];
            seen[nums[i]] = true;
            int imb = 0;
            for (int j = i+1; j < n; ++j) {
                int x = nums[j];
                if (!seen[x]) {
                    if (!seen[x-1] && !seen[x+1]) imb++;
                    if (seen[x-1] && seen[x+1]) imb--;
                }
                seen[x] = true;
                res += imb;
            }
        }
        return res;
    }
}

Kotlin

class Solution {
    fun sumImbalanceNumbers(nums: IntArray): Int {
        val n = nums.size
        var res = 0
        for (i in 0 until n) {
            val seen = BooleanArray(n+2)
            seen[nums[i]] = true
            var imb = 0
            for (j in i+1 until n) {
                val x = nums[j]
                if (!seen[x]) {
                    if (!seen[x-1] && !seen[x+1]) imb++
                    if (seen[x-1] && seen[x+1]) imb--
                }
                seen[x] = true
                res += imb
            }
        }
        return res
    }
}

Python

class Solution:
    def sumImbalanceNumbers(self, nums: list[int]) -> int:
        n = len(nums)
        res = 0
        for i in range(n):
            seen = [False] * (n+2)
            seen[nums[i]] = True
            imb = 0
            for j in range(i+1, n):
                x = nums[j]
                if not seen[x]:
                    if not seen[x-1] and not seen[x+1]:
                        imb += 1
                    if seen[x-1] and seen[x+1]:
                        imb -= 1
                seen[x] = True
                res += imb
        return res

Rust

impl Solution {
    pub fn sum_imbalance_numbers(nums: Vec<i32>) -> i32 {
        let n = nums.len();
        let mut res = 0;
        for i in 0..n {
            let mut seen = vec![false; n+2];
            seen[nums[i] as usize] = true;
            let mut imb = 0;
            for j in i+1..n {
                let x = nums[j] as usize;
                if !seen[x] {
                    if !seen[x-1] && !seen[x+1] { imb += 1; }
                    if seen[x-1] && seen[x+1] { imb -= 1; }
                }
                seen[x] = true;
                res += imb;
            }
        }
        res
    }
}

TypeScript

function sumImbalanceNumbers(nums: number[]): number {
    const n = nums.length;
    let res = 0;
    for (let i = 0; i < n; ++i) {
        const seen = new Array(n+2).fill(false);
        seen[nums[i]] = true;
        let imb = 0;
        for (let j = i+1; j < n; ++j) {
            const x = nums[j];
            if (!seen[x]) {
                if (!seen[x-1] && !seen[x+1]) imb++;
                if (seen[x-1] && seen[x+1]) imb--;
            }
            seen[x] = true;
            res += imb;
        }
    }
    return res;
}

---

Complexity

• ⏰ Time complexity: O(n^2)
• 🧺 Space complexity: O(n)