CSES Solutions Subarray Distinct Values (original) (raw)

Last Updated : 23 Jul, 2025

Given an array of N integersarr[] 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 tokeep 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:

C++ `

#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;

}

Java

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!
}

}

JavaScript

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

Python3

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))

`

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