LeetCode-in-Java

2763. Sum of Imbalance Numbers of All Subarrays

Hard

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:

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:

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:

The imbalance number of all other subarrays is 0. Hence, the sum of imbalance numbers of all the subarrays of nums is 8.

Constraints:

Solution

import java.util.Arrays;

public class Solution {
    public int sumImbalanceNumbers(int[] nums) {
        int n = nums.length;
        int s = 0;
        int[] left = new int[n];
        int[] seen = new int[n + 2];
        Arrays.fill(seen, -1);
        for (int i = 0; i < n; i++) {
            left[i] = Math.max(seen[nums[i]], seen[nums[i] + 1]);
            seen[nums[i]] = i;
        }
        Arrays.fill(seen, n);
        for (int i = n - 1; i >= 0; i--) {
            s += (seen[nums[i] + 1] - i) * (i - left[i]);
            seen[nums[i]] = i;
        }
        return s - (n + 1) * n / 2;
    }
}