Longest Subarray With At Most K Distinct Elements

Problem

Find the length of the longest subarray that contains at most k distinct elements.

Input:  nums = [1,2,1,2,3], k = 2
Output: 4   ([1,2,1,2])

Python Solution

def longest_subarray_k_distinct(nums, k):
    counts = {}
    left = 0
    best = 0
    for right, num in enumerate(nums):
        counts[num] = counts.get(num, 0) + 1
        while len(counts) > k:
            counts[nums[left]] -= 1
            if counts[nums[left]] == 0:
                del counts[nums[left]]
            left += 1
        best = max(best, right - left + 1)
    return best

C++ Solution

int longestSubarrayKDistinct(vector<int>& nums, int k) {
    unordered_map<int, int> counts;
    int left = 0, best = 0;
    for (int right = 0; right < nums.size(); right++) {
        counts[nums[right]]++;
        while (counts.size() > k) {
            counts[nums[left]]--;
            if (counts[nums[left]] == 0) counts.erase(nums[left]);
            left++;
        }
        best = max(best, right - left + 1);
    }
    return best;
}

Java Solution

class Solution {
    public int longestSubarrayKDistinct(int[] nums, int k) {
        Map<Integer, Integer> counts = new HashMap<>();
        int left = 0, best = 0;
        for (int right = 0; right < nums.length; right++) {
            counts.merge(nums[right], 1, Integer::sum);
            while (counts.size() > k) {
                int leftNum = nums[left];
                counts.merge(leftNum, -1, Integer::sum);
                if (counts.get(leftNum) == 0) counts.remove(leftNum);
                left++;
            }
            best = Math.max(best, right - left + 1);
        }
        return best;
    }
}

Complexity

Time: O(n) — the left pointer only moves forward, and each element enters and leaves the map a constant number of times. Space: O(k) — the map holds at most k+1 distinct elements at any point.