Subarrays with K Different Integers

Problem

Given an integer array nums and an integer k, return _the number ofgood subarrays of _nums.

A good array is an array where the number of different integers in that array is exactly k.

• For example, [1,2,3,1,2] has 3 different integers: 1, 2, and 3.

A subarray is a contiguous part of an array.

Examples

Example 1

Input: nums = [1,2,1,2,3], k = 2
Output: 7
Explanation: Subarrays formed with exactly 2 different integers: [1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], [1,2,1,2]

Example 2

Input: nums = [1,2,1,3,4], k = 3
Output: 3
Explanation: Subarrays formed with exactly 3 different integers: [1,2,1,3], [2,1,3], [1,3,4].

Constraints

• 1 <= nums.length <= 2 * 10^4
• 1 <= nums[i], k <= nums.length

Solution

Method 1 - Sliding Window (at most K) Trick

Intuition

The number of subarrays with exactly K different integers is equal to the number with at most K minus the number with at most K-1. We can use a sliding window to count subarrays with at most K different integers efficiently.

Approach

1. Implement a helper function that returns the number of subarrays with at most K different integers using a sliding window and a hash map.
2. The answer is helper(nums, K) - helper(nums, K-1).

Code

C++

#include <vector>
#include <unordered_map>
using namespace std;
class Solution {
    int atMostK(vector<int>& nums, int K) {
        unordered_map<int, int> count;
        int res = 0, left = 0;
        for (int right = 0; right < nums.size(); ++right) {
            if (++count[nums[right]] == 1) --K;
            while (K < 0) {
                if (--count[nums[left++]] == 0) ++K;
            }
            res += right - left + 1;
        }
        return res;
    }
public:
    int subarraysWithKDistinct(vector<int>& nums, int K) {
        return atMostK(nums, K) - atMostK(nums, K-1);
    }
};

Go

func subarraysWithKDistinct(nums []int, K int) int {
    return atMostK(nums, K) - atMostK(nums, K-1)
}
func atMostK(nums []int, K int) int {
    count := map[int]int{}
    res, left := 0, 0
    for right, v := range nums {
        count[v]++
        if count[v] == 1 { K-- }
        for K < 0 {
            count[nums[left]]--
            if count[nums[left]] == 0 { K++ }
            left++
        }
        res += right - left + 1
    }
    return res
}

Java

import java.util.*;
class Solution {
    public int subarraysWithKDistinct(int[] nums, int K) {
        return atMostK(nums, K) - atMostK(nums, K-1);
    }
    private int atMostK(int[] nums, int K) {
        Map<Integer, Integer> count = new HashMap<>();
        int res = 0, left = 0;
        for (int right = 0; right < nums.length; ++right) {
            count.put(nums[right], count.getOrDefault(nums[right], 0) + 1);
            if (count.get(nums[right]) == 1) --K;
            while (K < 0) {
                count.put(nums[left], count.get(nums[left]) - 1);
                if (count.get(nums[left]) == 0) ++K;
                left++;
            }
            res += right - left + 1;
        }
        return res;
    }
}

Kotlin

fun subarraysWithKDistinct(nums: IntArray, K: Int): Int {
    fun atMostK(K: Int): Int {
        val count = mutableMapOf<Int, Int>()
        var res = 0; var left = 0
        for (right in nums.indices) {
            count[nums[right]] = count.getOrDefault(nums[right], 0) + 1
            if (count[nums[right]] == 1) K--
            while (K < 0) {
                count[nums[left]] = count[nums[left]]!! - 1
                if (count[nums[left]] == 0) K++
                left++
            }
            res += right - left + 1
        }
        return res
    }
    return atMostK(K) - atMostK(K-1)
}

Python

def subarraysWithKDistinct(nums, K):
    def atMostK(K):
        count = {}
        res = left = 0
        for right, v in enumerate(nums):
            count[v] = count.get(v, 0) + 1
            if count[v] == 1:
                K -= 1
            while K < 0:
                count[nums[left]] -= 1
                if count[nums[left]] == 0:
                    K += 1
                left += 1
            res += right - left + 1
        return res
    return atMostK(K) - atMostK(K-1)

Rust

use std::collections::HashMap;
pub fn subarrays_with_k_distinct(nums: Vec<i32>, k: i32) -> i32 {
    fn at_most_k(nums: &Vec<i32>, mut k: i32) -> i32 {
        let mut count = HashMap::new();
        let mut res = 0;
        let mut left = 0;
        for right in 0..nums.len() {
            *count.entry(nums[right]).or_insert(0) += 1;
            if *count.get(&nums[right]).unwrap() == 1 { k -= 1; }
            while k < 0 {
                *count.get_mut(&nums[left]).unwrap() -= 1;
                if *count.get(&nums[left]).unwrap() == 0 { k += 1; }
                left += 1;
            }
            res += (right - left + 1) as i32;
        }
        res
    }
    at_most_k(&nums, k) - at_most_k(&nums, k-1)
}

TypeScript

function subarraysWithKDistinct(nums: number[], K: number): number {
    function atMostK(K: number): number {
        const count: Record<number, number> = {};
        let res = 0, left = 0;
        for (let right = 0; right < nums.length; right++) {
            count[nums[right]] = (count[nums[right]] || 0) + 1;
            if (count[nums[right]] === 1) K--;
            while (K < 0) {
                count[nums[left]]--;
                if (count[nums[left]] === 0) K++;
                left++;
            }
            res += right - left + 1;
        }
        return res;
    }
    return atMostK(K) - atMostK(K-1);
}

Complexity

• ⏰ Time complexity: O(n)
• 🧺 Space complexity: O(n)