Subarrays with K Different Integers
Problem
Given an integer array nums and an integer k, return the number of subarrays with exactly k distinct integers.
- 1 ≤ nums.length ≤ 2 * 10⁴
- 1 ≤ nums[i], k ≤ nums.length
Example
nums = [1,2,1,2,3], k = 27The subarrays with exactly 2 distinct integers are: [1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], and [1,2,1,2]. The algorithm uses a helper function to count subarrays with at most k distinct integers and subtracts the count of subarrays with at most k-1 distinct integers to get the exact count. For k=2, atMost(2) counts all subarrays with up to 2 distinct integers, and atMost(1) counts those with up to 1 distinct integer. Their difference yields the count of subarrays with exactly 2 distinct integers.
Approach
Straightforward Solution
A brute-force approach would enumerate all subarrays and count distinct integers in each, resulting in O(n^2) or worse time complexity, which is too slow for large inputs.
Core Observation
The number of subarrays with exactly k distinct integers equals the number of subarrays with at most k distinct integers minus the number of subarrays with at most k-1 distinct integers. This transforms the problem into counting subarrays with at most k distinct integers, which can be efficiently done using a sliding window.
Path to Optimal
The key insight is to use the sliding window technique to count subarrays with at most k distinct integers in O(n) time. By maintaining a window with at most k distinct integers and expanding the right pointer while adjusting the left pointer to maintain the constraint, the algorithm counts all valid subarrays ending at each right index. Applying this twice for k and k-1 and subtracting yields the count of subarrays with exactly k distinct integers.
Optimal Approach
Implement a helper function atMost(k) that uses a sliding window to count subarrays with at most k distinct integers. The main function returns atMost(k) - atMost(k - 1). This approach runs in O(n) time with O(n) space for the frequency map.
Time
O(n)
Each element is visited at most twice: once when expanding the right pointer and once when moving the left pointer. The frequency map operations are O(1) on average, resulting in an overall linear time complexity.
Space
O(n)
The frequency map stores counts of elements currently in the sliding window. In the worst case, it can hold all distinct elements in nums, which is O(n). This auxiliary space is necessary to track distinct counts efficiently.
Pattern Spotlight
Sliding Window (Counting Subarrays with Constraints)
When counting subarrays with an exact number of distinct elements, transform the problem into the difference of counts of subarrays with at most k and at most k-1 distinct elements, enabling efficient linear-time sliding window solutions.
Solution
| 1 | class Solution: |
| 2 | def subarraysWithKDistinct(self, nums, k): |
| 3 | |
| 4 | def atMost(k): |
| 5 | count = {} |
| 6 | left = 0 |
| 7 | res = 0 |
| 8 | |
| 9 | for right, n in enumerate(nums): |
| 10 | |
| 11 | count[n] = count.get(n, 0) + 1 |
| 12 | |
| 13 | while len(count) > k: |
| 14 | count[nums[left]] -= 1 |
| 15 | |
| 16 | if count[nums[left]] == 0: |
| 17 | del count[nums[left]] |
| 18 | |
| 19 | left += 1 |
| 20 | |
| 21 | res += right - left + 1 |
| 22 | |
| 23 | return res |
| 24 | |
| 25 | return atMost(k) - atMost(k - 1) |
Step-by-Step Solution
Count Subarrays with At Most K Distinct Integers Using Sliding Window
| 4 | def atMost(k): |
| 5 | count = {} |
| 6 | left = 0 |
| 7 | res = 0 |
| 9 | for right, n in enumerate(nums): |
| 11 | count[n] = count.get(n, 0) + 1 |
| 13 | while len(count) > k: |
| 14 | count[nums[left]] -= 1 |
| 16 | if count[nums[left]] == 0: |
| 17 | del count[nums[left]] |
| 19 | left += 1 |
| 21 | res += right - left + 1 |
| 23 | return res |
Objective
To count the number of subarrays with at most k distinct integers by expanding and contracting a sliding window.
Key Insight
The sliding window maintains a frequency map of elements within the window and adjusts its left boundary to ensure the window contains at most k distinct integers. For each right pointer position, the number of valid subarrays ending at that position is the window size (right - left + 1). Summing these counts yields the total number of subarrays with at most k distinct integers. This approach avoids enumerating all subarrays explicitly, achieving linear time complexity.
Interview Quick-Check
Core Logic
The sliding window expands the right pointer and contracts the left pointer to maintain at most k distinct integers, counting valid subarrays by window size at each step.
State & Boundaries
The while loop condition ensures the window never exceeds k distinct integers by removing elements from the left until the constraint is satisfied.
Common Pitfalls & Bugs
Forgetting to delete keys from the frequency map when their count drops to zero can cause incorrect distinct counts and infinite loops.
Complexity
Each element is added and removed from the frequency map at most once, ensuring O(n) time complexity.
Calculate Exact Count by Subtracting Counts of At Most K and K-1
| 25 | return atMost(k) - atMost(k - 1) |
Objective
To compute the number of subarrays with exactly k distinct integers by subtracting counts of subarrays with at most k-1 distinct integers from those with at most k.
Key Insight
The difference between the counts of subarrays with at most k and at most k-1 distinct integers isolates subarrays with exactly k distinct integers. This leverages the inclusion-exclusion principle and avoids direct enumeration of exact k distinct subarrays, which is more complex. The final return statement calls the helper function twice and returns their difference, completing the solution efficiently.
Interview Quick-Check
Core Logic
Subtracting atMost(k - 1) from atMost(k) isolates subarrays with exactly k distinct integers by removing those with fewer distinct integers.
Common Pitfalls & Bugs
Not handling the case when k is 0 or negative can cause errors; the problem constraints ensure k >= 1.
Line Analysis
This solution has 4 Critical lines interviewers watch for.
return atMost(k) - atMost(k - 1)Return the difference between counts of subarrays with at most k and at most k-1 distinct integers.
This subtraction isolates the count of subarrays with exactly k distinct integers, leveraging the helper function's results and the inclusion-exclusion principle.
while len(count) > k:While the window contains more than k distinct integers, shrink it from the left.
This loop enforces the constraint of at most k distinct integers by moving the left pointer forward and updating frequencies accordingly.
del count[nums[left]]Delete the element from the frequency map so the distinct count remains accurate.
Removing elements whose count becomes zero ensures the dictionary size accurately represents the number of distinct elements in the window, which is required for the len(count) constraint check.
res += right - left + 1Add the number of valid subarrays ending at the current right pointer to the result.
The count of valid subarrays ending at right is the window size (right - left + 1), as all subarrays starting between left and right are valid.
count = {}Initialize a dictionary to track the frequency of elements in the current window.
The frequency map is essential to efficiently track the number of distinct elements in the sliding window and update counts as the window moves.
count[n] = count.get(n, 0) + 1Increment the frequency count for the current element at the right pointer.
Updating the frequency map reflects the inclusion of the new element in the window, which may increase the distinct count.
count[nums[left]] -= 1Decrement the frequency count of the element at the left pointer.
Reducing the count reflects removing the leftmost element from the window, potentially decreasing the distinct count.
if count[nums[left]] == 0:Check whether the element removed from the left side of the window now has zero frequency.
After decrementing the count of the element leaving the window, we check whether its frequency has dropped to zero. If it has, that element is no longer present in the sliding window.
def atMost(k):Define a helper function to count subarrays with at most k distinct integers.
Encapsulating the sliding window logic in a helper function allows reuse for both k and k-1, enabling the calculation of exact counts via subtraction.
for right, n in enumerate(nums):Iterate over the array with the right pointer expanding the sliding window.
The right pointer explores all possible window end positions, enabling the algorithm to consider every subarray ending at each index.
left += 1Move the left pointer forward to shrink the window.
Advancing the left pointer reduces the window size, restoring the at most k distinct integers constraint.
return resReturn the total count of subarrays with at most k distinct integers.
Returning the accumulated result completes the helper function, enabling its use in the main function to compute exact counts.
left = 0Initialize the left pointer of the sliding window.
The left pointer marks the start of the current window and moves forward to maintain the distinct count constraint.
res = 0Initialize the result accumulator for counting valid subarrays.
This variable accumulates the total count of subarrays with at most k distinct integers as the window expands.
Test Your Understanding
Why does subtracting the count of subarrays with at most k-1 distinct integers from the count of subarrays with at most k distinct integers yield the count of subarrays with exactly k distinct integers?
Because every subarray with exactly k distinct integers is counted in atMost(k) but not in atMost(k-1). Subarrays with fewer than k distinct integers are counted in both, so subtracting removes those, leaving only subarrays with exactly k distinct integers.
Related Problems
Sliding Window pattern
Don't just read it. Drill it.
Reconstruct Subarrays with K Different Integers from memory until it sticks. AlgoDrill blanks out key lines and makes you fill them back in, step by step.
Unlock the Subarrays with K Different Integers drill