CSES Solutions - Subarray Distinct Values

Given an array of N integers arr[] and an integer K, your task is to calculate the number of subarrays that have at most K distinct values.

Examples:

Input: N = 5, K = 2, arr[] = {1, 2, 1, 2, 3}
Output: 10
Explanation: There are 12 subarrays with at most 2 distinct values: {1}, {2}, {1}, {2}, {3}, {1, 2}, {2, 1}, {1, 2}, {2, 3}, {1, 2, 1}, {2, 1, 2}, {1, 2, 1, 2}.

Input: N = 4, K = 1, arr[] = {2, 2, 2, 2}
Output: 4
Explanation: There are 4 subarrays with at most 1 distinct values: {2}, {2}, {2}, {2}.

Approach To solve the problem, follow the below idea:

We maintain a Sliding Window and a hashmap to track the count of elements in the window. As we move the window's right boundary, we update the hashmap. If the number of distinct elements in the window exceeds k, we shrink the window from the left until it contains at most k distinct elements. We count the subarrays for each valid window position and add them to the result.

Step-by-step algorithm:

• Initialize two pointers left = 0 and right = 0 to mark the starting and ending point of the sliding window and another variable result = 0 to store the number of arrays with at most K distinct values.
• Maintain a variable distinct_count to keep track of the number of distinct elements in the current window.
• Also maintain a map, say frequency to store the frequency of elements in the current window.
• Iterate over the array till the right pointer of the window does not exceed the end of the array.
• Move the right pointer till the number of distinct elements in the current window <= K.
• Shrink the window from the left side if the number of distinct elements exceeds K.
• Decrement the frequency of the leftmost element and move the left pointer to the right until the number of distinct elements becomes <= K.
• While doing this, we accumulate the count of valid subarrays by adding the length of each subarray to result.
• After reaching the end of the array, we print the final result.

Below is the implementation of the above algorithm:

12

Time Complexity: O(N), where N is the size of array arr[].
Auxiliary Space: O(N)