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:

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

Output

12

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