Count Subarrays With At Most K Distinct Elements

Given an array arr[] of integers and a positive integer k, the goal is to count the total number of subarrays that contain at most k distinct (unique) elements.

Examples:

Input: arr[] = [1, 2, 2, 3], k = 2 
Output: 9
Explanation: Subarrays with at most 2 distinct elements are: [1], [2], [2], [3], [1, 2], [2, 2], [2, 3], [1, 2, 2] and [2, 2, 3].

Input: arr[] = [1, 1, 1], k = 1
Output: 6
Explanation: Subarrays with at most 1 distinct element are: [1], [1], [1], [1, 1], [1, 1] and [1, 1, 1].

Input: arr[] = [1, 2, 1, 1, 3, 3, 4, 2, 1], k = 2
Output: 24
Explanation: There are 24 subarrays with at most 2 distinct elements.

[Naive Approach] Exploring all subarrays – O(n^2) Time and O(n) Space

The idea is to iterate over all possible subarrays while keeping the count of distinct integers using a hash set. For every subarray, check if the size of hash set is less than or equal to k. If the size of hash set is less than or equal to k, increment the result by 1. After iterating over all subarrays, return the result.

24

[Expected Approach] Using Sliding Window Technique - O(n) Time and O(k) Space

The idea is to use Sliding Window Technique by using two pointers, say left and right to mark the start and end boundary of the window. Initialize both the pointers to the first element. Expand the right boundary until distinct element count exceeds k, then shrink the left boundary until the distinct element count becomes <= k. While expanding the window, for each right boundary keep counting the subarrays as (right – left + 1).

The intuition behind (right – left + 1) is that it counts all possible subarrays that end at right and start at any index between left and right. These are valid subarrays because they all contain at most k distinct elements.

Working:

24