Skip to content

Subarrays with K Different Integers⚓︎

Link

Description⚓︎

Given an integer array nums and an integer k, return the number of good 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.

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⚓︎

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
private:
    int numOfMostKDistinct(vector<int>& nums, int k) {
        int n = nums.size();
        int res = 0, kinds = 0;
        vector<int> cnts(2e4+1, 0);

        for (int left = 0, right = 0; right < n; right++) {
            if (++cnts[nums[right]] == 1) kinds++;

            while (kinds > k) {
                if (--cnts[nums[left++]] == 0) {
                    kinds--;
                }
            }
            res += right - left + 1;
        }
        return res;
    }

public:
    int subarraysWithKDistinct(vector<int>& nums, int k) {
        return numOfMostKDistinct(nums, k) - numOfMostKDistinct(nums, k - 1);
    }
};

Code with comments:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
private:
    // Helper function to calculate the number of subarrays with at most K distinct integers
    int numOfMostKDistinct(vector<int>& nums, int k) {
        int n = nums.size(); // Size of the input array
        int res = 0; // Result variable to store the number of valid subarrays
        int kinds = 0; // Variable to keep track of the number of distinct integers in the current window
        vector<int> cnts(2e4+1, 0); // Count array to keep track of the frequency of each integer within the current window

        // Two pointers technique: `left` and `right` define the current window boundaries
        for (int left = 0, right = 0; right < n; right++) {
            // If this is the first occurrence of nums[right], increment `kinds`
            if (++cnts[nums[right]] == 1) kinds++;

            // If the number of distinct integers exceeds `k`, shrink the window from the left
            while (kinds > k) {
                if (--cnts[nums[left]] == 0) { // If removing nums[left] makes its count 0, decrement `kinds`
                    kinds--;
                }
                left++; // Move the left boundary of the window to the right
            }

            // The number of subarrays ending at `right` with at most `K` distinct integers
            // is `right - left + 1`. Accumulate this to `res`.
            res += right - left + 1;
        }
        return res;
    }

public:
    // Main function to find the number of subarrays with exactly K distinct integers
    int subarraysWithKDistinct(vector<int>& nums, int k) {
        // The difference between the number of subarrays with at most `K` distinct integers
        // and the number of subarrays with at most `K-1` distinct integers gives the required answer.
        return numOfMostKDistinct(nums, k) - numOfMostKDistinct(nums, k - 1);
    }
};

Explanation:

The solution relies on a sliding window approach to count subarrays with at most K distinct integers, encapsulated in the numOfMostKDistinct function. The main idea is to maintain a window [left, right] that satisfies the condition of having at most K distinct integers. This is achieved by incrementing right to expand the window and incrementing left to shrink the window whenever the condition is violated (i.e., when the number of distinct integers exceeds K).

For each position of right, the solution counts how many subarrays ending at right meet the criteria by calculating the difference right - left + 1. This effectively counts all valid subarrays that include nums[right].

To find the number of subarrays with exactly K distinct integers, the solution exploits a clever trick: it calculates the number of subarrays with at most K distinct integers and subtracts from it the number of subarrays with at most K-1 distinct integers. This subtraction removes all subarrays that have fewer than K distinct integers, leaving only those with exactly K distinct integers.

This technique simplifies the problem by converting it from counting subarrays with "exactly" to "at most" K distinct integers, which is easier to manage using a sliding window approach.