LeetCode in Kotlin

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

class Solution {
    fun sumImbalanceNumbers(nums: IntArray): Int {
        val n = nums.size
        var ans = 0
        for (i in 0 until n) {
            val s: MutableSet<Int> = HashSet()
            var curr = 0
            for (j in i until n) {
                val x = nums[j]
                if (s.contains(x)) {
                    // do nothing
                } else if (s.contains(x - 1) && s.contains(x + 1)) {
                    curr--
                } else if (!s.contains(x - 1) && !s.contains(x + 1) && s.isNotEmpty()) {
                    curr++
                }
                s.add(x)
                ans += curr
            }
        }
        return ans
    }
}