🔗 CSES-2428 Distinct Values Subarrays II

Problem Statement

題目簡述

給定一個長度為 $n$ 的整數陣列，求有多少個子陣列，其包含的不同元素個數至多為 $k$ 。

Constraints

約束條件

$1 \le k \le n \le 2 \cdot 10^5$
$1 \le x_i \le 10^9$

思路：滑動窗口（雙指標）

1. 從暴力到單調性

最直觀的想法是枚舉所有可能的子陣列 $[l, r]$ ，並統計其中有多少個不同的元素。但這樣做的時間複雜度是 $\mathcal{O}(n^2)$ ，對於 $n \le 2 \cdot 10^5$ 的數據範圍顯然會逾時。

我們需要尋找優化的突破口。

關鍵觀察：區間的單調性

假設一個子陣列 $[l, r]$ 內的不同元素個數至多為 $k$ 。
如果我們固定右端點 $r$ ，並將左端點向右移動到 $l'$ （即 $l < l' \le r$ ），那麼子陣列 $[l', r]$ 內的不同元素個數只可能減少，不可能增加。

這意味著：如果 $[l, r]$ 是合法的，那麼 $[l', r]$ 也必然是合法的。

反過來說，對於每個固定的右端點 $r$ ，滿足條件的左端點會形成一個連續的區間。我們只需要找到最左邊的合法左端點 $l_{\min}$ ，那麼所有起點在 $[l_{\min}, r]$ 範圍內、終點為 $r$ 的子陣列都是合法的。

這樣的合法子陣列數量，恰好就是當前合法窗口的長度：

r - l_{\min} + 1

2. 滑動窗口雙指標模板

既然左端點具有單調性（隨著右端點 $r$ 向右移動，最左合法左端點 $l_{\min}$ 也只會向右移動，不會向左回退），我們就可以使用 雙指標（滑動窗口） 在 $\mathcal{O}(n)$ 時間內解決此問題。

核心套路：枚舉右，維護左

這是滑動窗口的經典思維模型：

枚舉右端點：讓右端點一步步向右擴展，將新元素加入窗口，並更新窗口內元素的出現次數。
收縮左端點：如果此時窗口內的不同元素個數超過了 $k$ ，我們就必須不斷將左端點向右移動，直到窗口內的不同元素個數重新回到 $\le k$ 的範圍。
累加貢獻：此時以當前右端點結尾的合法子陣列個數就是 $r - l + 1$ ，將其累加到答案中。

雜湊表維護的細節

在收縮左端點時，我們會減少左端點對應元素的計數。
若用雜湊表的大小來表示窗口內不同元素個數，當某個元素的計數降為 $0$ 時，必須將其從雜湊表中徹底刪除（del）。
否則，雜湊表的大小（鍵的個數）將無法正確反映窗口內「不同元素」的實際數量。

另一種寫法是手動維護目前窗口內的不同元素數量：加入一個原本不存在於窗口的元素時，數量加一；移除一個計數降為 $0$ 的元素時，數量減一。這樣即使不刪除雜湊表中的鍵，也能正確判斷窗口是否合法。

如果值域很小，或元素只包含小寫字母、大小寫字母等固定集合，也可以把雜湊表換成普通陣列計數，並搭配手動維護的不同元素數量。

複雜度分析

時間複雜度： $\mathcal{O}(n)$ 。右端點和左端點都最多移動 $n$ 次，雜湊表的操作（插入、查詢、刪除）平均為 $\mathcal{O}(1)$ ，因此總時間複雜度為線性。
空間複雜度： $\mathcal{O}(k)$ 。雜湊表中最多同時存在 $k + 1$ 個不同的元素。

類題

滑動窗口(Sliding Window)

340. Longest Substring with At Most K Distinct Characters（越短越合法，求最大長度）
992. Subarrays with K Different Integers（雙指標進階：恰好 $K$ 個轉化為至多 $K$ 個減去至多 $K-1$ 個）

Code

在實作中，我們使用一個雜湊表（在 Python 中為 defaultdict）來記錄當前窗口內每個元素的出現次數。雜湊表的長度（鍵的個數）即為窗口內不同元素的個數。

from collections import defaultdict


def solve():
    n, k = map(int, input().split())
    A = list(map(int, input().split()))
    assert len(A) == n

    ans = left = 0
    cnt = defaultdict(int)
    for right, x in enumerate(A):
        cnt[x] += 1
        # 當窗口內不同元素個數超過 k 時，收縮左端點
        while len(cnt) > k:
            y = A[left]
            cnt[y] -= 1
            if cnt[y] == 0:
                del cnt[y]
            left += 1
        # 累加當前右端點結尾的合法子陣列數量
        ans += right - left + 1
    print(ans)


if __name__ == "__main__":
    solve()