Sum of Imbalance Numbers of All Subarrays
HardPrompt
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, andsarr[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-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 <= 10001 <= nums[i] <= nums.length
Approaches
4 approaches with complexity analysis and trade-offs.
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.
Algorithm
- Initialize a variable
totalImbalanceto 0. - Generate all possible subarrays of
nums. This can be done using two nested loops, with the outer loop for the starting indexiand the inner loop for the ending indexj. - For each subarray
nums[i...j]: a. Create a temporary copy of the subarray. b. Sort the temporary copy. Let's call itsarr. c. Calculate the imbalance number forsarr. Iterate from the first element to the second-to-last element ofsarrand increment a counter ifsarr[k+1] - sarr[k] > 1. d. Add this imbalance number tototalImbalance. - After iterating through all subarrays, return
totalImbalance.
Walkthrough
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.
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; }}Complexity
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.
Trade-offs
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.
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.
Algorithm
- Initialize
totalImbalanceto 0. - Iterate through each possible starting index
iof a subarray. - For each
i, initialize a sorted data structure (like aTreeSet) and a variablecurrentImbalancefor the current subarray being extended. - Iterate with a second loop for the ending index
jfromiton-1. - In the inner loop, add
nums[j]to the sorted set. As you add the element, updatecurrentImbalancebased on how the new element affects the gaps with its neighbors in the sorted set. - If
nums[j]is a new unique element, find its predecessor (prev) and successor (next) in the set. The introduction ofnums[j]might remove one old gap (next - prev) and introduce two new gaps (nums[j] - prev,next - nums[j]). UpdatecurrentImbalanceaccordingly. - Add
currentImbalancetototalImbalancein each step of the inner loop. - Return
totalImbalance.
Walkthrough
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.
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; }}Complexity
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.
Trade-offs
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.
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.
Algorithm
- 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. - Initialize
totalImbalanceto 0. - Iterate through each starting index
ifrom0ton-1. - For each
i, initialize a set/boolean arrayseento track unique numbers and a counteradjacentPairsto 0. - Iterate with an inner loop for the ending index
jfromiton-1. - For each element
x = nums[j]: a. Ifxis already inseen, the imbalance doesn't change. b. Ifxis new, add it toseen. Check ifx-1andx+1are inseen. Ifseen.contains(x-1), incrementadjacentPairs. Ifseen.contains(x+1), incrementadjacentPairs. c. Calculate the current subarray's imbalance:currentImbalance = seen.size() - 1 - adjacentPairs. (Handle the edge case whereseen.size() <= 1, imbalance is 0). d. AddcurrentImbalancetototalImbalance. - Return
totalImbalance.
Walkthrough
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.
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; }}Complexity
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`.
Trade-offs
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.
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.
Algorithm
- The total sum can be expressed as:
Sum(imbalance(S)) = Sum(|unique(S)| - 1) - Sum(count_adjacent_pairs(S))over all subarraysS. - Calculate Term 1:
Sum(|unique(S)| - 1)a. This isSum(|unique(S)|) - Sum(1).Sum(1)is the total number of subarrays,n(n+1)/2. b.Sum(|unique(S)|)equalsSum over x (number of subarrays containing x). c. For each valuevfrom1ton, find all its indices. The number of subarrays not containingvcan be calculated from the gaps between its occurrences. Subtract this from the total number of subarrays to get the number of subarrays containingv. d. Sum these counts for allvto getSum(|unique(S)|). - Calculate Term 2:
Sum(count_adjacent_pairs(S))a. This equalsSum 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), andN(A or B)can all be calculated using the same gap-based method from step 2c, by considering the indices ofv,v+1, and their union, respectively. - Final Result: The answer is
Term1 - Term2.
Walkthrough
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).
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; }}Complexity
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.
Trade-offs
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.
Solutions
Solution
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; }}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; }};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 ansVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.