如何计算数组滑动窗口内不同元素的数量并优化解决超时问题

原有代码问题 #

你原来的实现每次都对长度为B的窗口切片后转set统计不同元素，时间复杂度为O(NB)*，当数组长度和窗口大小较大时会触发超时。

优化思路 #

使用频率字典记录当前窗口内各元素的出现次数，搭配全局变量记录当前窗口的不同元素数量，每次窗口滑动仅需处理移出的左边界元素和新增的右边界元素，避免重复统计整个窗口。具体步骤如下：

• 先统计第一个长度为B的窗口内各元素的频率，计算得到第一个窗口的不同元素数量，加入结果数组
• 从第B个元素开始遍历数组，每次窗口右移一位：
  • 处理左边界移出元素：将对应元素的频率减1，若减到0说明该元素已完全移出当前窗口，不同元素数量减1
  • 处理右边界新增元素：若该元素当前频率为0说明是首次进入当前窗口，不同元素数量加1，再将对应元素的频率加1
  • 将当前窗口的不同元素数量加入结果数组

优化后代码 #

from collections import defaultdict

def count_distinct_sliding_window(A, B):
    freq = defaultdict(int)
    res = []
    distinct_cnt = 0
    # 初始化第一个窗口
    for i in range(B):
        num = A[i]
        if freq[num] == 0:
            distinct_cnt += 1
        freq[num] += 1
    res.append(distinct_cnt)
    # 滑动窗口遍历
    for i in range(B, len(A)):
        # 移出左边界元素
        left_num = A[i - B]
        freq[left_num] -= 1
        if freq[left_num] == 0:
            distinct_cnt -= 1
        # 加入右边界新元素
        right_num = A[i]
        if freq[right_num] == 0:
            distinct_cnt += 1
        freq[right_num] += 1
        res.append(distinct_cnt)
    return res

# 测试示例
A = [1, 2, 1, 3, 4, 3]
B = 3
print(count_distinct_sliding_window(A, B)) # 输出 [2, 3, 3, 2]

复杂度说明 #

• 时间复杂度：O(N)，仅需遍历数组两次，所有哈希表操作均为均摊O(1)
• 空间复杂度：O(K)，K为数组中不同元素的总个数，最坏情况为O(N)

内容的提问来源于stack exchange，提问作者Maws