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:
#include <bits/stdc++.h>
#define ll long long
using namespace std;
ll solve(ll* arr, ll N, ll K)
{
// left and right pointers to mark the start and end of
// the sliding window
ll left = 0, right = 0;
// Variable to count how many different numbers we have
// in the window
ll distinct_count = 0;
// Variable to store the final result
ll result = 0;
// Map to keep track of how many times each number
// appears in the window
unordered_map<ll, ll> frequency;
// Slide the window till the window till the right
// pointer does not reach the end of the array
while (right < N) {
// Check if the current number is new or if its
// count is zero
if (frequency.find(arr[right]) == frequency.end()
|| frequency[arr[right]] == 0)
distinct_count++;
// Update the count of the current number
frequency[arr[right]]++;
// If there are more than K distinct numbers, shrink
// the window from the left
while (distinct_count > K) {
// Decrease the count of the number going out of
// the window
frequency[arr[left]]--;
// If its count becomes zero, then there will be
// one less distinct number in the window
if (frequency[arr[left]] == 0)
distinct_count--;
// Move the left pointer to the right to shrink
// the window
left++;
}
// Calculate the number of subarrays that end at the
// current position
result += right - left + 1;
// Move the right edge of the window to the right to
// expand it
right++;
}
// Return the result
return result;
}
int main()
{
// Sample Input
ll N = 5, K = 2;
ll arr[N] = { 1, 2, 1, 2, 3 };
cout << solve(arr, N, K) << "\n";
// We're done!
return 0;
}
import java.util.HashMap;
public class Main {
static long solve(long[] arr, long N, long K) {
// left and right pointers to mark the start and end of
// the sliding window
long left = 0, right = 0;
// Variable to count how many different numbers we have
// in the window
long distinct_count = 0;
// Variable to store the final result
long result = 0;
// Map to keep track of how many times each number
// appears in the window
HashMap<Long, Long> frequency = new HashMap<>();
// Slide the window till the window till the right
// pointer does not reach the end of the array
while (right < N) {
// Check if the current number is new or if its
// count is zero
if (!frequency.containsKey(arr[(int)right]) || frequency.get(arr[(int)right]) == 0)
distinct_count++;
// Update the count of the current number
frequency.put(arr[(int)right], frequency.getOrDefault(arr[(int)right], 0L) + 1);
// If there are more than K distinct numbers, shrink
// the window from the left
while (distinct_count > K) {
// Decrease the count of the number going out of
// the window
frequency.put(arr[(int)left], frequency.get(arr[(int)left]) - 1);
// If its count becomes zero, then there will be
// one less distinct number in the window
if (frequency.get(arr[(int)left]) == 0)
distinct_count--;
// Move the left pointer to the right to shrink
// the window
left++;
}
// Calculate the number of subarrays that end at the
// current position
result += right - left + 1;
// Move the right edge of the window to the right to
// expand it
right++;
}
// Return the result
return result;
}
public static void main(String[] args) {
// Sample Input
long N = 5, K = 2;
long[] arr = { 1, 2, 1, 2, 3 };
System.out.println(solve(arr, N, K));
// We're done!
}
}
function solve(arr, N, K) {
// left and right pointers to mark the start and end of
// the sliding window
let left = 0, right = 0;
// Variable to count how many different numbers we have
// in the window
let distinctCount = 0;
// Variable to store the final result
let result = 0;
// Map to keep track of how many times each number
// appears in the window
let frequency = {};
// Slide the window till the right pointer does not reach the end of the array
while (right < N) {
// Check if the current number is new or if its count is zero
if (!frequency[arr[right]] || frequency[arr[right]] === 0)
distinctCount++;
// Update the count of the current number
frequency[arr[right]] = (frequency[arr[right]] || 0) + 1;
// If there are more than K distinct numbers, shrink the window from the left
while (distinctCount > K) {
// Decrease the count of the number going out of the window
frequency[arr[left]]--;
// If its count becomes zero, then there will be one less distinct number in the window
if (frequency[arr[left]] === 0)
distinctCount--;
// Move the left pointer to the right to shrink the window
left++;
}
// Calculate the number of subarrays that end at the current position
result += right - left + 1;
// Move the right edge of the window to the right to expand it
right++;
}
return result;
}
let N = 5, K = 2;
let arr = [1, 2, 1, 2, 3];
console.log(solve(arr, N, K));
// This code is contributed by Ayush Mishra
from collections import defaultdict
def solve(arr, N, K):
# left and right pointers to mark the start and end of
# the sliding window
left = 0
right = 0
# Variable to count how many different numbers we have
# in the window
distinct_count = 0
# Variable to store the final result
result = 0
# Dictionary to keep track of how many times each number
# appears in the window
frequency = defaultdict(int)
# Slide the window till the window till the right
# pointer does not reach the end of the array
while right < N:
# Check if the current number is new or if its
# count is zero
if frequency[arr[right]] == 0:
distinct_count += 1
# Update the count of the current number
frequency[arr[right]] += 1
# If there are more than K distinct numbers, shrink
# the window from the left
while distinct_count > K:
# Decrease the count of the number going out of
# the window
frequency[arr[left]] -= 1
# If its count becomes zero, then there will be
# one less distinct number in the window
if frequency[arr[left]] == 0:
distinct_count -= 1
# Move the left pointer to the right to shrink
# the window
left += 1
# Calculate the number of subarrays that end at the
# current position
result += right - left + 1
# Move the right edge of the window to the right to
# expand it
right += 1
# Return the result
return result
N = 5
K = 2
arr = [1, 2, 1, 2, 3]
print(solve(arr, N, K))
Output
12
Time Complexity: O(N), where N is the size of array arr[].
Auxiliary Space: O(N)
Contact Us