2762. Continuous Subarrays

Problem

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 <= i₁, i₂ <= j, 0 <= |nums[i₁] - nums[i₂]| <= 2.

Return the total number of continuous subarrays.

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

https://leetcode.com/problems/continuous-subarrays/

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⁵
1 <= nums[i] <= 10⁹

Test Cases

1 2	class Solution: def continuousSubarrays(self, nums: List[int]) -> int:

solution_test.py

import pytest

from solution import Solution


@pytest.mark.parametrize('nums, expected', [
    ([5,4,2,4], 8),
    ([1,2,3], 6),
])
class Test:
    def test_solution(self, nums, expected):
        sol = Solution()
        assert sol.continuousSubarrays(nums) == expected

Thoughts

用滑动窗口扫描数组，跟 3258. Count Substrings That Satisfy K-Constraint I 几乎一致。但是在 problem 3258 中每次是统计 i 开头的子串数，这里改成统计 j 结尾的子数组数，这样最会留个尾数需要加。

移窗过程中的操作类似于 76. Minimum Window Substring，用哈希表记录当前窗口中不同数字的个数。可以用常数时间判定当前窗口是否符合要求，当窗口移动时也用常数时间更新哈希表。

时间复杂度 O(n)，空间复杂度 O(1)。

Code

solution.py

from collections import defaultdict


class Solution:
    def continuousSubarrays(self, nums: list[int]) -> int:
        n = len(nums)
        cnt = 0
        counts = defaultdict(int) # counts[v]: the number of `v`s in current window.
        i = j = 0
        counts[nums[j]] = 1
        while True:
            if max(counts) - min(counts) <= 2:
                j += 1
                cnt += j - i
                if j == n: break

                val = nums[j]
                counts[val] += 1
            else:
                val = nums[i]
                counts[val] -= 1
                if counts[val] == 0: del counts[val]
                i += 1

        return cnt