Skip to content

2763. Sum of Imbalance Numbers of All Subarrays 👍

Approach 1: $O(n^2)$

  • Time: $O(n^2)$
  • Space: $O(n)$

Approach 2: $O(n)$

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Solution {
 public:
  // If sorted(nums)[i + 1] - sorted(nums)[i] > 1, then there's a gap. Instead
  // of determining the number of gaps in each subarray, let's find out how many
  // subarrays contain each gap.
  int sumImbalanceNumbers(vector<int>& nums) {
    const int n = nums.size();
    int ans = 0;
    // Note that to avoid double counting, only `left` needs to check nums[i].
    // This adjustment ensures that i represents the position of the leftmost
    // element of nums[i] within the subarray.

    // left[i] := the maximum index l s.t. nums[l] = nums[i] or nums[i] + 1
    vector<int> left(n);
    // right[i] := the minimum index r s.t. nums[r] = nums[i]
    vector<int> right(n);

    vector<int> numToIndex(n + 2, -1);
    for (int i = 0; i < n; ++i) {
      left[i] = max(numToIndex[nums[i]], numToIndex[nums[i] + 1]);
      numToIndex[nums[i]] = i;
    }

    fill(numToIndex.begin(), numToIndex.end(), n);
    for (int i = n - 1; i >= 0; --i) {
      right[i] = numToIndex[nums[i] + 1];
      numToIndex[nums[i]] = i;
    }

    // The gap above nums[i] persists until encountering nums[i] or nums[i] + 1.
    // Consider subarrays nums[l..r] with l <= i <= r, where l in [left[i], i]
    // and r in [i, right[i] - 1]. There are (i - left[i]) * (right[i] - i)
    // subarrays satisfying this condition.
    for (int i = 0; i < n; ++i)
      ans += (i - left[i]) * (right[i] - i);

    // Subtract n * (n + 1) / 2 to account for the overcounting of elements
    // initially assumed to have a gap. This adjustment is necessary as the
    // maximum element of every subarray does not have a gap.
    return ans - n * (n + 1) / 2;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Solution {
  // If sorted(nums)[i + 1] - sorted(nums)[i] > 1, then there's a gap. Instead
  // of determining the number of gaps in each subarray, let's find out how many
  // subarrays contain each gap.
  public int sumImbalanceNumbers(int[] nums) {
    final int n = nums.length;
    int ans = 0;
    // Note that to avoid double counting, only `left` needs to check nums[i].
    // This adjustment ensures that i represents the position of the leftmost
    // element of nums[i] within the subarray.

    // left[i] := the maximum index l s.t. nums[l] = nums[i] or nums[i] + 1
    int[] left = new int[n];
    // right[i] := the minimum index r s.t. nums[r] = nums[i]
    int[] right = new int[n];
    int[] numToIndex = new int[n + 2];

    Arrays.fill(numToIndex, -1);
    for (int i = 0; i < n; ++i) {
      left[i] = Math.max(numToIndex[nums[i]], numToIndex[nums[i] + 1]);
      numToIndex[nums[i]] = i;
    }

    Arrays.fill(numToIndex, n);
    for (int i = n - 1; i >= 0; --i) {
      right[i] = numToIndex[nums[i] + 1];
      numToIndex[nums[i]] = i;
    }

    // The gap above nums[i] persists until encountering nums[i] or nums[i] + 1.
    // Consider subarrays nums[l..r] with l <= i <= r, where l in [left[i], i]
    // and r in [i, right[i] - 1]. There are (i - left[i]) * (right[i] - i)
    // subarrays satisfying this condition.
    for (int i = 0; i < n; ++i)
      ans += (i - left[i]) * (right[i] - i);

    // Subtract n * (n + 1) / 2 to account for the overcounting of elements
    // initially assumed to have a gap. This adjustment is necessary as the
    // maximum element of every subarray does not have a gap.
    return ans - n * (n + 1) / 2;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class Solution:
  # If sorted(nums)[i + 1] - sorted(nums)[i] > 1, then there's a gap. Instead
  # of determining the number of gaps in each subarray, let's find out how many
  # subarrays contain each gap.
  def sumImbalanceNumbers(self, nums: List[int]) -> int:
    n = len(nums)
    # Note that to avoid double counting, only `left` needs to check nums[i].
    # This adjustment ensures that i represents the position of the leftmost
    # element of nums[i] within the subarray.

    # left[i] := the maximum index l s.t. nums[l] = nums[i] or nums[i] + 1
    left = [0] * n
    # right[i] := the minimum index r s.t. nums[r] = nums[i]
    right = [0] * n

    numToIndex = [-1] * (n + 2)
    for i, num in enumerate(nums):
      left[i] = max(numToIndex[num], numToIndex[num + 1])
      numToIndex[num] = i

    numToIndex = [n] * (n + 2)
    for i in range(n - 1, -1, -1):
      right[i] = numToIndex[nums[i] + 1]
      numToIndex[nums[i]] = i

    # The gap above nums[i] persists until encountering nums[i] or nums[i] + 1.
    # Consider subarrays nums[l..r] with l <= i <= r, where l in [left[i], i]
    # and r in [i, right[i] - 1]. There are (i - left[i]) * (right[i] - i)
    # subarrays satisfying this condition.
    #
    # Subtract n * (n + 1) / 2 to account for the overcounting of elements
    # initially assumed to have a gap. This adjustment is necessary as the
    # maximum element of every subarray does not have a gap.
    return sum((i - left[i]) * (right[i] - i)
               for i in range(n)) - n * (n + 1) // 2