Count Subarrays With At Most K Distinct Elements

Last Updated : 23 Jul, 2025

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.

Table of Content

[Naive Approach] Exploring all subarrays – O(n^2) Time and O(n) Space
[Expected Approach] Using Sliding Window Technique - O(n) Time and O(k) Space

[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.

C++

#include <iostream>
#include <unordered_set>
#include <vector>
using namespace std;

int countAtMostK(vector<int> &arr, int k) {
    int n = arr.size();
    int res = 0;

    for (int i = 0; i < n; i++) {

        // hash set to count distinct elements
        unordered_set<int> st;

        for (int j = i; j < n; j++) {
            st.insert(arr[j]);

            // if count of distinct elements > k, then
            // don't extend the subarray further
            if (st.size() > k)
                break;
                
            // no. of distinct element is less than or equal k 
            res += 1;
        }
        
    }
    return res;
}

int main() {
    vector<int> arr = {1, 2, 1, 1, 3, 3, 4, 2, 1};
    int k = 2;
    cout << countAtMostK(arr, k);
}

Java

import java.util.HashSet;
import java.util.Set;

class GfG {
    static int countAtMostK(int[] arr, int k) {
        int n = arr.length;
        int res = 0;

        for (int i = 0; i < n; i++) {

            // hash set to count distinct elements
            Set<Integer> st = new HashSet<>();

            for (int j = i; j < n; j++) {
                st.add(arr[j]);

                // if count of distinct elements > k, then
                // don't extend the subarray further
                if (st.size() > k)
                    break;
                    
                // no. of distinct element is less than or equal k 
                res += 1;
            }
        }
        return res;
    }

    public static void main(String[] args) {
        int[] arr = {1, 2, 1, 1, 3, 3, 4, 2, 1};
        int k = 2;
        System.out.println(countAtMostK(arr, k));
    }
}

Python

def countAtMostK(arr, k):
    n = len(arr)
    res = 0

    for i in range(n):

        # hash set to count distinct elements
        st = set()

        for j in range(i, n):
            st.add(arr[j])

            # If count of distinct elements > k, then
            # don't extend the subarray further
            if len(st) > k:
                break
            
            # no. of distinct element is less than or equal k 
            res += 1
    return res

if __name__ == "__main__":
	arr = [1, 2, 1, 1, 3, 3, 4, 2, 1]
	k = 2
	print(countAtMostK(arr, k))

using System;
using System.Collections.Generic;

class GfG {
    static int countAtMostK(int[] arr, int k) {
        int n = arr.Length;
        int res = 0;

        for (int i = 0; i < n; i++) {
          
            // Hash set to count distinct elements
            HashSet<int> st = new HashSet<int>();

            for (int j = i; j < n; j++) {
                st.Add(arr[j]);

                // If count of distinct elements > k, then
                // don't extend the subarray further
                if (st.Count > k)
                    break;
                    
                // no. of distinct element is less than or equal k 
                res += 1;
            }
        }
        return res;
    }

    static void Main() {
        int[] arr = new int[] {1, 2, 1, 1, 3, 3, 4, 2, 1};
        int k = 2;
        Console.WriteLine(countAtMostK(arr, k));
    }
}

JavaScript

function countAtMostK(arr, k) {
    let n = arr.length;
    let res = 0;

    for (let i = 0; i < n; i++) {

        // Set to count distinct elements
        let st = new Set();

        for (let j = i; j < n; j++) {
            st.add(arr[j]);

            // If count of distinct elements > k, then
            // don't extend the subarray further
            if (st.size > k)
                break;
            
            // no. of distinct element is less than or equal k 
            res += 1;
        }
    }
    return res;
}

// Driver Code
let arr = [1, 2, 1, 1, 3, 3, 4, 2, 1];
let k = 2;
console.log(countAtMostK(arr, k));

Output

[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:

C++

#include <iostream>
#include <unordered_map>
#include <vector>
using namespace std;

int countAtMostK(vector<int> &arr, int k) {
    int n = arr.size();
    int res = 0;

    // pointers to mark the left and right boundary
    int left = 0, right = 0;

    // frequency map 
    unordered_map<int, int> freq;
    while (right < n) {
        freq[arr[right]] += 1;

        // if this is a new element in the window,
        // decrement k by 1
        if (freq[arr[right]] == 1)
            k -= 1;

        // shrink the window until distinct element
        // count becomes <= k
        while (k < 0) {
            freq[arr[left]] -= 1;
            if (freq[arr[left]] == 0)
                k += 1;
            left += 1;
        }

        // count of subarrays ending at "right" 
        // and having atmost k elements
        res += (right - left + 1);
        right += 1;
    }
    return res;
}

int main() {
    vector<int> arr = {1, 2, 1, 1, 3, 3, 4, 2, 1};
    int k = 2;
    cout << countAtMostK(arr, k);
}

Java

import java.util.HashMap;
import java.util.Map;

class GfG {

    static int countAtMostK(int[] arr, int k) {
        int n = arr.length;
        int res = 0;

        // pointers to mark the left and right boundary
        int left = 0, right = 0;

        // frequency map
        Map<Integer, Integer> freq = new HashMap<>();
        while (right < n) {
            freq.put(arr[right], freq.getOrDefault(arr[right], 0) + 1);

            // If this is a new element in the window, 
            // decrement k by 1
            if (freq.get(arr[right]) == 1)
                k -= 1;

            // shrink the window until distinct element
            // count becomes <= k
            while (k < 0) {
                freq.put(arr[left], freq.get(arr[left]) - 1);
                if (freq.get(arr[left]) == 0)
                    k += 1;
                left += 1;
            }

            // count of subarrays ending at "right" 
            // and having atmost k elements
            res += (right - left + 1);
            right += 1;
        }
        return res;
    }

    public static void main(String[] args) {
        int[] arr = {1, 2, 1, 1, 3, 3, 4, 2, 1};
        int k = 2;
        System.out.println(countAtMostK(arr, k));
    }
}

Python

from collections import defaultdict

def countAtMostK(arr, k):
    n = len(arr)
    res = 0

    # pointers to mark the left and right boundary
    left, right = 0, 0

    # frequency map
    freq = defaultdict(int)
    while right < n:
        freq[arr[right]] += 1

        # if this is a new element in the window,
        # decrement k by 1
        if freq[arr[right]] == 1:
            k -= 1

        # shrink the window until distinct element
        # count becomes <= k
        while k < 0:
            freq[arr[left]] -= 1
            if freq[arr[left]] == 0:
                k += 1
            left += 1

        # count of subarrays ending at "right" 
        # and having atmost k elements
        res += (right - left + 1)
        right += 1
    return res


if __name__ == "__main__":
    arr = [1, 2, 1, 1, 3, 3, 4, 2, 1]
    k = 2
    print(countAtMostK(arr, k))

using System;
using System.Collections.Generic;

class GfG {
    static int countAtMostK(int[] arr, int k) {
        int n = arr.Length;
        int res = 0;

        // pointers to mark the left and right boundary
        int left = 0, right = 0;

        // frequency map
        Dictionary<int, int> freq = new Dictionary<int, int>();
        while (right < n) {
            if (freq.ContainsKey(arr[right]))
                freq[arr[right]] += 1;
            else
                freq[arr[right]] = 1;

            // if this is a new element in the window,
            // decrement k by 1
            if (freq[arr[right]] == 1)
                k -= 1;

            // shrink the window until distinct element
            // count becomes <= k
            while (k < 0) {
                freq[arr[left]] -= 1;
                if (freq[arr[left]] == 0)
                    k += 1;
                left += 1;
            }

            // count of subarrays ending at "right" 
            // and having atmost k elements
            res += (right - left + 1);
            right += 1;
        }
        return res;
    }

    static void Main(string[] args) {
        int[] arr = { 1, 2, 1, 1, 3, 3, 4, 2, 1 };
        int k = 2;
        Console.WriteLine(countAtMostK(arr, k));
    }
}

JavaScript

function countAtMostK(arr, k) {
    let n = arr.length;
    let res = 0;

    // pointers to mark the left and right boundary
    let left = 0, right = 0;

    // frequency map
    let freq = new Map();
    while (right < n) {
        freq.set(arr[right], (freq.get(arr[right]) || 0) + 1);

        // if this is a new element in the window, 
        // decrement k by 1
        if (freq.get(arr[right]) === 1)
            k -= 1;

        // shrink the window until distinct element
        // count becomes <= k
        while (k < 0) {
            freq.set(arr[left], freq.get(arr[left]) - 1);
            if (freq.get(arr[left]) === 0)
                k += 1;
            left += 1;
        }

        // count of subarrays ending at "right" 
        // and having atmost k elements
        res += (right - left + 1);
        right += 1;
    }
    return res;
}

// Driver Code
let arr = [1, 2, 1, 1, 3, 3, 4, 2, 1];
let k = 2;
console.log(countAtMostK(arr, k));