# Sum of Imbalance Numbers of All Subarrays
**Difficulty:** HARD
[External](https://leetcode.com/problems/sum-of-imbalance-numbers-of-all-subarrays)
Canonical: https://scaleengineer.com/dsa/problems/sum-of-imbalance-numbers-of-all-subarrays
**Data structures:** Array, Hash Table, Ordered Set
---
## 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 _the **sum of imbalance numbers** of all its **subarrays**_.

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

**Example 1:**

**Input:** nums = [2,3,1,4]
**Output:** 3
**Explanation:** There are 3 subarrays with non-zeroimbalance 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`

# Approaches
## Brute Force with Re-sorting
The most straightforward approach is to simulate the process directly. We can generate every single subarray of the input array `nums`. For each of these subarrays, we calculate its imbalance number as defined in the problem statement and add it to a running total. The imbalance number calculation involves sorting the subarray and then counting the gaps greater than 1 between adjacent elements.
**Time:** O(n³ log n). There are O(n²) subarrays. For each subarray of length up to n, we spend O(n log n) to sort it. This results in a cubic time complexity, which is too slow for n=1000. · **Space:** O(n), for storing the temporary subarray.
**Pros:** Simple to understand and implement.; Directly follows the problem definition.
**Cons:** Extremely inefficient and will time out for the given constraints.; Repeatedly sorts subarrays, which involves a lot of redundant computation.
### Explanation
This method iterates through all possible start and end points to define a subarray. For each subarray, a new list is created, sorted, and then scanned to find its imbalance number. The sum of these imbalance numbers over all subarrays gives the final answer.

For example, with `nums = [2,3,1,4]`, we would first consider the subarray `[2]`. Sorted, it's `[2]`. Imbalance is 0. Then `[2,3]`. Sorted, it's `[2,3]`. Imbalance is 0. Then `[2,3,1]`. Sorted, it's `[1,2,3]`. Imbalance is 0. This process continues for all `n*(n+1)/2` subarrays.

```java
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class Solution {
    public int sumImbalanceNumbers(int[] nums) {
        int n = nums.length;
        int totalImbalance = 0;

        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                // Create the subarray
                List<Integer> sub = new ArrayList<>();
                for (int k = i; k <= j; k++) {
                    sub.add(nums[k]);
                }

                // Calculate imbalance for the subarray
                if (sub.size() <= 1) {
                    continue;
                }

                Collections.sort(sub);

                int currentImbalance = 0;
                for (int k = 0; k < sub.size() - 1; k++) {
                    if (sub.get(k + 1) - sub.get(k) > 1) {
                        currentImbalance++;
                    }
                }
                totalImbalance += currentImbalance;
            }
        }
        return totalImbalance;
    }
}
```
### Algorithm
1. Initialize a variable `totalImbalance` to 0.
2. Generate all possible subarrays of `nums`. This can be done using two nested loops, with the outer loop for the starting index `i` and the inner loop for the ending index `j`.
3. For each subarray `nums[i...j]`:
    a. Create a temporary copy of the subarray.
    b. Sort the temporary copy. Let's call it `sarr`.
    c. Calculate the imbalance number for `sarr`. Iterate from the first element to the second-to-last element of `sarr` and increment a counter if `sarr[k+1] - sarr[k] > 1`.
    d. Add this imbalance number to `totalImbalance`.
4. After iterating through all subarrays, return `totalImbalance`.

## Iterate Subarrays with Sorted Set
This approach improves upon the brute-force method by avoiding re-sorting the entire subarray from scratch. We iterate through all subarrays, but as we extend a subarray by one element (from `nums[i...j-1]` to `nums[i...j]`), we maintain a sorted data structure of its unique elements, like a `TreeSet`. When adding a new element, we can efficiently update the imbalance count in logarithmic time by checking its immediate neighbors in the sorted structure.
**Time:** O(n² log n). The two nested loops give O(n²), and each operation inside the inner loop (insertion and search in `TreeSet`) takes O(log k) where k is the size of the subarray, which is at most n. · **Space:** O(n), for the `TreeSet` which can store up to `n` elements.
**Pros:** More efficient than the naive brute-force approach.; Avoids redundant sorting operations.
**Cons:** The logarithmic time complexity for set operations in the inner loop makes it slower than the optimal O(n²) approach.
### Explanation
For each starting position `i`, we build subarrays `nums[i...j]` by progressively increasing `j`. A `TreeSet` is used to keep track of the unique elements of the current subarray in sorted order. When a new element `x = nums[j]` is added, we find its sorted neighbors, `prev` (the largest element in the set smaller than `x`) and `next` (the smallest element larger than `x`). The change in imbalance is calculated by observing how the gaps are affected. For instance, if a gap `next - prev > 1` existed, it's now replaced by two new gaps, `x - prev` and `next - x`. We adjust the imbalance count based on whether these new gaps are greater than 1.

```java
import java.util.TreeSet;

class Solution {
    public int sumImbalanceNumbers(int[] nums) {
        int n = nums.length;
        int totalImbalance = 0;

        for (int i = 0; i < n; i++) {
            TreeSet<Integer> s = new TreeSet<>();
            int currentImbalance = 0;
            for (int j = i; j < n; j++) {
                int x = nums[j];
                if (s.contains(x)) {
                    // If element already exists, unique elements don't change, so imbalance is the same.
                    totalImbalance += currentImbalance;
                    continue;
                }

                Integer prev = s.lower(x);
                Integer next = s.higher(x);

                if (prev != null && next != null) {
                    if (next - prev > 1) currentImbalance--;
                    if (x - prev > 1) currentImbalance++;
                    if (next - x > 1) currentImbalance++;
                } else if (prev != null) {
                    if (x - prev > 1) currentImbalance++;
                } else if (next != null) {
                    if (next - x > 1) currentImbalance++;
                }
                
                s.add(x);
                totalImbalance += currentImbalance;
            }
        }
        return totalImbalance;
    }
}
```
### Algorithm
1. Initialize `totalImbalance` to 0.
2. Iterate through each possible starting index `i` of a subarray.
3. For each `i`, initialize a sorted data structure (like a `TreeSet`) and a variable `currentImbalance` for the current subarray being extended.
4. Iterate with a second loop for the ending index `j` from `i` to `n-1`.
5. In the inner loop, add `nums[j]` to the sorted set. As you add the element, update `currentImbalance` based on how the new element affects the gaps with its neighbors in the sorted set.
6. If `nums[j]` is a new unique element, find its predecessor (`prev`) and successor (`next`) in the set. The introduction of `nums[j]` might remove one old gap (`next - prev`) and introduce two new gaps (`nums[j] - prev`, `next - nums[j]`). Update `currentImbalance` accordingly.
7. Add `currentImbalance` to `totalImbalance` in each step of the inner loop.
8. Return `totalImbalance`.

## Iterate Subarrays with Hashing
This approach refines the previous one by optimizing the inner loop. Instead of a `TreeSet`, we can use a `HashSet` or a boolean array (since element values are bounded by `n`) for O(1) average time lookups. The key insight is a different formula for the imbalance number: `imbalance = (number of unique elements) - 1 - (count of adjacent pairs)`. An adjacent pair is a pair of numbers `(v, v+1)` that are both present in the subarray. As we extend the subarray `nums[i...j]`, we can efficiently update the count of unique elements and adjacent pairs in O(1) time.
**Time:** O(n²). Two nested loops give a quadratic time complexity. The operations inside the inner loop are O(1). · **Space:** O(n), for the `seen` boolean array or `HashSet`.
**Pros:** Efficient enough to pass the given constraints.; Improves upon the O(n² log n) approach by using a hash set or boolean array for O(1) updates.
**Cons:** Still has a quadratic time complexity, which might be slow for very large N, but is fine for the given constraints.
### Explanation
We iterate through all subarrays with a nested loop. For each starting index `i`, we maintain a set of numbers seen so far in the subarray starting at `i`. As we extend the subarray to the right by including `nums[j]`, we update our counts. If `nums[j]` is a new number, we increment the unique count. We also check if its neighbors (`nums[j]-1` and `nums[j]+1`) are already in our set. If they are, it means we've formed new adjacent pairs, so we update our `adjacentPairs` count. The imbalance for the current subarray `nums[i...j]` is then calculated and added to the total sum.

Since `1 <= nums[i] <= nums.length`, we can use a boolean array for the `seen` set for maximum efficiency.

```java
import java.util.HashSet;
import java.util.Set;

class Solution {
    public int sumImbalanceNumbers(int[] nums) {
        int n = nums.length;
        int totalImbalance = 0;

        for (int i = 0; i < n; i++) {
            // Using a boolean array is faster than HashSet due to constraints
            boolean[] seen = new boolean[n + 2];
            int uniqueCount = 0;
            int adjacentPairs = 0;

            for (int j = i; j < n; j++) {
                int x = nums[j];

                if (!seen[x]) {
                    uniqueCount++;
                    if (seen[x - 1]) {
                        adjacentPairs++;
                    }
                    if (seen[x + 1]) {
                        adjacentPairs++;
                    }
                    seen[x] = true;
                }

                if (uniqueCount <= 1) {
                    totalImbalance += 0;
                } else {
                    totalImbalance += (uniqueCount - 1 - adjacentPairs);
                }
            }
        }
        return totalImbalance;
    }
}
```
### Algorithm
1. The imbalance of a set can be calculated as `(number of unique elements) - 1 - (number of adjacent pairs)`. An adjacent pair is `(x, x+1)` where both are in the set.
2. Initialize `totalImbalance` to 0.
3. Iterate through each starting index `i` from `0` to `n-1`.
4. For each `i`, initialize a set/boolean array `seen` to track unique numbers and a counter `adjacentPairs` to 0.
5. Iterate with an inner loop for the ending index `j` from `i` to `n-1`.
6. For each element `x = nums[j]`:
    a. If `x` is already in `seen`, the imbalance doesn't change.
    b. If `x` is new, add it to `seen`. Check if `x-1` and `x+1` are in `seen`. If `seen.contains(x-1)`, increment `adjacentPairs`. If `seen.contains(x+1)`, increment `adjacentPairs`.
    c. Calculate the current subarray's imbalance: `currentImbalance = seen.size() - 1 - adjacentPairs`. (Handle the edge case where `seen.size() <= 1`, imbalance is 0).
    d. Add `currentImbalance` to `totalImbalance`.
7. Return `totalImbalance`.

## Combinatorial Counting
The most optimal solution uses a combinatorial approach to count the total imbalance in linear time. Instead of iterating through subarrays, we change the order of summation. The total imbalance is the sum of imbalances of all subarrays. We can rewrite the imbalance formula and calculate the total sum of each of its components across all subarrays.

The total imbalance is `Σ (|unique(S)| - 1) - Σ (count of adjacent pairs in S)`. We can calculate these two terms separately in O(n) time. The first term involves counting, for each value `x`, how many subarrays contain it. The second term involves counting, for each pair `(x, x+1)`, how many subarrays contain both. These counts can be derived efficiently by first finding all indices of each number and then using a combinatorial argument based on the gaps between these indices.
**Time:** O(n). Pre-calculating positions takes O(n). The main loops run up to `n` times. The work inside the loop for a value `v` depends on the number of its occurrences. Since the sum of occurrences of all numbers is `n`, the total work across all loops is linear. · **Space:** O(n), to store the positions of each number.
**Pros:** Most efficient solution with linear time complexity.; Scales well to much larger inputs than specified in the constraints.
**Cons:** The logic is complex and harder to implement correctly compared to simpler approaches.; Requires careful handling of calculations to avoid overflow and off-by-one errors.
### Explanation
This method avoids iterating through subarrays altogether. It relies on a mathematical reformulation of the problem.

First, we calculate `Term1 = Σ (|unique(S)| - 1)`. This is `(Σ |unique(S)|) - (total number of subarrays)`. We can find `Σ |unique(S)|` by summing up, for each distinct number `v`, the number of subarrays that contain `v`. This can be computed in O(n) total by pre-calculating the positions of each number.

Second, we calculate `Term2 = Σ (count of adjacent pairs in S)`. This is equivalent to `Σ_{v=1}^{n-1} (number of subarrays containing both v and v+1)`. For each pair `(v, v+1)`, we find the number of subarrays containing both using the inclusion-exclusion principle. The number of subarrays containing a set of values can be found by looking at the indices of these values and counting the subarrays that exist in the 'gaps' between them.

This approach requires helper functions to calculate the number of subarrays containing a given set of values based on their indices. The overall time complexity is dominated by iterating through the values from 1 to n, with calculations for each step being proportional to the number of occurrences, leading to a total of O(n).

```java
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class Solution {
    public int sumImbalanceNumbers(int[] nums) {
        int n = nums.length;
        List<Integer>[] pos = new ArrayList[n + 1];
        for (int i = 0; i <= n; i++) {
            pos[i] = new ArrayList<>();
        }
        for (int i = 0; i < n; i++) {
            pos[nums[i]].add(i);
        }

        long totalSubarrays = (long) n * (n + 1) / 2;
        long sumUniqueSizes = 0;
        for (int v = 1; v <= n; v++) {
            if (!pos[v].isEmpty()) {
                sumUniqueSizes += countSubarraysWithValues(n, pos[v]);
            }
        }

        long term1 = sumUniqueSizes - totalSubarrays;

        long totalGoodPairs = 0;
        for (int v = 1; v < n; v++) {
            if (pos[v].isEmpty() || pos[v + 1].isEmpty()) {
                continue;
            }
            long countWithV = countSubarraysWithValues(n, pos[v]);
            long countWithV1 = countSubarraysWithValues(n, pos[v + 1]);

            List<Integer> union = new ArrayList<>(pos[v]);
            union.addAll(pos[v + 1]);
            Collections.sort(union);
            long countWithVOrV1 = countSubarraysWithValues(n, union);

            totalGoodPairs += countWithV + countWithV1 - countWithVOrV1;
        }

        return (int) (term1 - totalGoodPairs);
    }

    private long countSubarraysWithValues(int n, List<Integer> indices) {
        long totalSubarrays = (long) n * (n + 1) / 2;
        long subarraysWithout = 0;
        int lastIndex = -1;
        for (int index : indices) {
            long gap = index - lastIndex - 1;
            subarraysWithout += gap * (gap + 1) / 2;
            lastIndex = index;
        }
        long gap = n - lastIndex - 1;
        subarraysWithout += gap * (gap + 1) / 2;
        return totalSubarrays - subarraysWithout;
    }
}
```
### Algorithm
1. The total sum can be expressed as: `Sum(imbalance(S)) = Sum(|unique(S)| - 1) - Sum(count_adjacent_pairs(S))` over all subarrays `S`.
2. **Calculate Term 1: `Sum(|unique(S)| - 1)`**
   a. This is `Sum(|unique(S)|) - Sum(1)`. `Sum(1)` is the total number of subarrays, `n(n+1)/2`.
   b. `Sum(|unique(S)|)` equals `Sum over x (number of subarrays containing x)`.
   c. For each value `v` from `1` to `n`, find all its indices. The number of subarrays *not* containing `v` can be calculated from the gaps between its occurrences. Subtract this from the total number of subarrays to get the number of subarrays containing `v`.
   d. Sum these counts for all `v` to get `Sum(|unique(S)|)`.
3. **Calculate Term 2: `Sum(count_adjacent_pairs(S))`**
   a. This equals `Sum over v (number of subarrays containing both v and v+1)`.
   b. For each pair `(v, v+1)`, we can find the number of subarrays containing both using the principle of inclusion-exclusion: `N(A and B) = N(A) + N(B) - N(A or B)`.
   c. `N(A)`, `N(B)`, and `N(A or B)` can all be calculated using the same gap-based method from step 2c, by considering the indices of `v`, `v+1`, and their union, respectively.
4. **Final Result:** The answer is `Term1 - Term2`.

# Solutions
### Java

```java
class Solution {
public
  int sumImbalanceNumbers(int[] nums) {
    int n = nums.length;
    int ans = 0;
    for (int i = 0; i < n; ++i) {
      TreeMap<Integer, Integer> tm = new TreeMap<>();
      int cnt = 0;
      for (int j = i; j < n; ++j) {
        Integer k = tm.ceilingKey(nums[j]);
        if (k != null && k - nums[j] > 1) {
          ++cnt;
        }
        Integer h = tm.floorKey(nums[j]);
        if (h != null && nums[j] - h > 1) {
          ++cnt;
        }
        if (h != null && k != null && k - h > 1) {
          --cnt;
        }
        tm.merge(nums[j], 1, Integer : : sum);
        ans += cnt;
      }
    }
    return ans;
  }
}

```

### CPP

```cpp
class Solution {
public:
  int sumImbalanceNumbers(vector<int> &nums) {
    int n = nums.size();
    int ans = 0;
    for (int i = 0; i < n; ++i) {
      multiset<int> s;
      int cnt = 0;
      for (int j = i; j < n; ++j) {
        auto it = s.lower_bound(nums[j]);
        if (it != s.end() && *it - nums[j] > 1) {
          ++cnt;
        }
        if (it != s.begin() && nums[j] - *prev(it) > 1) {
          ++cnt;
        }
        if (it != s.end() && it != s.begin() && *it - *prev(it) > 1) {
          --cnt;
        }
        s.insert(nums[j]);
        ans += cnt;
      }
    }
    return ans;
  }
};

```

### Python

```python
from sortedcontainers import SortedList class Solution : def sumImbalanceNumbers ( self , nums : List [ int ]) -> int : n = len ( nums ) ans = 0 for i in range ( n ): sl = SortedList () cnt = 0 for j in range ( i , n ): k = sl . bisect_left ( nums [ j ]) h = k - 1 if h >= 0 and nums [ j ] - sl [ h ] > 1 : cnt += 1 if k < len ( sl ) and sl [ k ] - nums [ j ] > 1 : cnt += 1 if h >= 0 and k < len ( sl ) and sl [ k ] - sl [ h ] > 1 : cnt -= 1 sl . add ( nums [ j ]) ans += cnt return ans
```