leetcode - 2762. Continuous Subarrays

Description

You are given a 0-indexed integer array nums. A subarray of nums is called continuous if:

Let i, i + 1, ..., j be the indices in the subarray. Then, for each pair of indices i <= i1, i2 <= j, 0 <= |nums[i1] - nums[i2]| <= 2.
Return the total number of continuous subarrays.

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

Example 1:

Input: nums = [5,4,2,4]
Output: 8
Explanation: 
Continuous subarray of size 1: [5], [4], [2], [4].
Continuous subarray of size 2: [5,4], [4,2], [2,4].
Continuous subarray of size 3: [4,2,4].
Thereare no subarrys of size 4.
Total continuous subarrays = 4 + 3 + 1 = 8.
It can be shown that there are no more continuous subarrays.

Example 2:

Input: nums = [1,2,3]
Output: 6
Explanation: 
Continuous subarray of size 1: [1], [2], [3].
Continuous subarray of size 2: [1,2], [2,3].
Continuous subarray of size 3: [1,2,3].
Total continuous subarrays = 3 + 2 + 1 = 6.

Constraints:

1 <= nums.length <= 10^5
1 <= nums[i] <= 10^9

Solution

Use a max-heap and a min-heap to keep track of all the element within a window, shrink when the difference between largest and smallest is larger than 2.

The number of subarray is the length of current window.

Time complexity: $o(n\log n)$
Space complexity: $o(n)$

Code

class Solution:
    def continuousSubarrays(self, nums: List[int]) -> int:
        left = 0
        res = 0
        max_heap, min_heap = [], []
        for right in range(len(nums)):
            # put right in
            heapq.heappush(max_heap, (-nums[right], right))
            heapq.heappush(min_heap, (nums[right], right))
            # shrink if needed
            while -max_heap[0][0] - min_heap[0][0] > 2:
                left = max(left, min(max_heap[0][1], min_heap[0][1]) + 1)
                while max_heap and max_heap[0][1] < left:
                    heapq.heappop(max_heap)
                while min_heap and min_heap[0][1] < left:
                    heapq.heappop(min_heap)
            res += right - left + 1
        return res