K 个不同整数的子数组

标签: 数组 哈希表 计数 滑动窗口

难度: Hard

给定一个正整数数组 nums和一个整数 k,返回 nums 中 「好子数组」 的数目。

如果 nums 的某个子数组中不同整数的个数恰好为 k,则称 nums 的这个连续、不一定不同的子数组为 好子数组 」

  • 例如,[1,2,3,1,2] 中有 3 个不同的整数:12,以及 3

子数组 是数组的 连续 部分。

示例 1:

输入:nums = [1,2,1,2,3], k = 2
输出:7
解释:恰好由 2 个不同整数组成的子数组:[1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], [1,2,1,2].

示例 2:

输入:nums = [1,2,1,3,4], k = 3
输出:3
解释:恰好由 3 个不同整数组成的子数组:[1,2,1,3], [2,1,3], [1,3,4].

提示:

  • 1 <= nums.length <= 2 * 104
  • 1 <= nums[i], k <= nums.length

Submission

运行时间: 96 ms

内存: 18.0 MB

class Solution:
    def subarraysWithKDistinct(self, nums: List[int], k: int) -> int:
        return self.atMostKDistinct(nums, k) - self.atMostKDistinct(nums, k-1)
    def atMostKDistinct(self, nums, k):
        n = len(nums)
        freq = [0] * (n+1)
        left, right = 0, 0
        cnt, res = 0, 0

        while right < n:
            if freq[nums[right]] == 0:
                cnt += 1
            freq[nums[right]] += 1
            right += 1

            while cnt > k:
                freq[nums[left]] -= 1
                if freq[nums[left]] == 0:
                    cnt -= 1
                left += 1
            res += right - left
        return res

Explain

这道题要求找出数组中恰好包含 k 个不同整数的子数组数目。题解采用了双指针技术和滑动窗口的方法,结合了两个具体的函数:计算最多有 k 个不同整数的子数组数目的函数 `atMostKDistinct`,以及通过计算 `atMostKDistinct(nums, k) - atMostKDistinct(nums, k-1)` 来得到恰好有 k 个不同整数的子数组数目。这种方法使用了频率数组来跟踪每个元素的出现次数,并动态调整窗口的大小来满足条件。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def subarraysWithKDistinct(self, nums: List[int], k: int) -> int:
        # 使用 atMostKDistinct 函数计算至多包含 k 个不同整数的子数组数目
        return self.atMostKDistinct(nums, k) - self.atMostKDistinct(nums, k-1)
    def atMostKDistinct(self, nums, k):
        n = len(nums)
        # 初始化频率数组,长度比 nums 多 1 以处理元素索引直接使用
        freq = [0] * (n+1)
        left, right = 0, 0
        cnt, res = 0, 0
        # 右指针负责扩展窗口
        while right < n:
            if freq[nums[right]] == 0:
                cnt += 1
            freq[nums[right]] += 1
            right += 1
            # 左指针负责缩小窗口直至窗口内不同整数数量不大于 k
            while cnt > k:
                freq[nums[left]] -= 1
                if freq[nums[left]] == 0:
                    cnt -= 1
                left += 1
            # 累计可能的子数组数目
            res += right - left
        return res

Explore

在`atMostKDistinct`函数中,当`nums[right]`的频率为0时增加计数器`cnt`是因为每当遇到一个新的整数(即此前未出现过的整数),我们需要增加不同整数的计数。频率为0表示此前该整数没有在当前考虑的窗口中出现过。增加`cnt`反映了窗口内不同整数的增加,这是为了确保我们可以正确追踪和控制窗口中包含的不同整数的数量,符合函数的目标——计算至多包含k个不同整数的子数组数目。

`right - left`代表的是当前有效窗口的长度。在滑动窗口方法中,每当我们扩展右指针`right`以包括一个新的元素,并且窗口内的不同整数的数量符合条件时(即不超过k个),窗口从左指针`left`到右指针`right`之间的任何一个子数组都是有效的。因此,`right - left`实际上表示以`nums[right]`结尾的、符合条件的子数组的数量,这些子数组的起始位置可以从`left`到`right`的任何位置。这一步骤通过累加这些可能的子数组数目来更新结果`res`。

在`atMostKDistinct`函数的实现中,右指针`right`总是先于左指针`left`移动,这确保了在任何调整左指针之前,右指针已经至少向前移动了一步。这种先移动右指针扩展窗口,再移动左指针缩小窗口的策略,确保了`left`指针永远不会超过`right`指针。此外,左指针的移动是在一个内部的循环中进行的,该循环只有在`cnt`(窗口中的不同整数数量)超过k时才会执行,并且循环的条件是`cnt > k`,这保证了左指针的移动是有条件和有限的,避免了越过右指针。

在减少`freq[nums[left]]`后检查是否变为0是因为我们需要准确追踪窗口内不同整数的数量。`freq[nums[left]]`表示整数`nums[left]`在窗口中的频率,当这个频率降至0时,意味着该整数已经完全从窗口中移除,因此不应再被计入窗口内不同整数的总数。于是,我们需要减少`cnt`以反映窗口内不同整数数量的减少。这是确保`cnt`准确反映当前窗口状态的关键步骤。