Scale Engineer

Subarrays with K Different Integers

Hard
#0946Time: O(N^2). There are two nested loops. The outer loop runs `N` times, and the inner loop runs up to `N` times. The operations inside the inner loop (adding to a `HashSet` and checking its size) take constant time on average.Space: O(k). The `HashSet` used to store distinct elements for a subarray will hold at most `k+1` elements before the inner loop breaks. In the worst case, `k` can be up to `N`, making it O(N).2 companies
Patterns
Sliding WindowCounting
Data structures
ArrayHash Table
Companies
Morgan StanleyRoblox

Prompt

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 * 104
  • 1 <= nums[i], k <= nums.length

Approaches

2 approaches with complexity analysis and trade-offs.

This approach involves checking every possible subarray of the input array nums. For each subarray, we count the number of distinct integers it contains. If this count is exactly k, we increment our result counter. While straightforward, this method is inefficient for large inputs due to its nested loop structure, leading to a quadratic time complexity.

Algorithm

  • Initialize a variable count to 0 to store the number of good subarrays.
  • Use a nested loop to generate all possible subarrays. The outer loop with index i determines the start of the subarray, and the inner loop with index j determines the end.
  • For each subarray starting at i, use a HashSet to keep track of the distinct elements encountered as j increases.
  • In the inner loop, for each element nums[j], add it to the HashSet.
  • After adding nums[j], check the size of the HashSet.
  • If distinctElements.size() == k, it means the current subarray nums[i...j] has exactly k distinct elements, so we increment count.
  • An optional optimization: If distinctElements.size() > k, we can break the inner loop because any further extension of this subarray will also have more than k distinct elements.
  • After iterating through all possible i and j, return the final count.

Walkthrough

The brute-force method systematically generates all subarrays and verifies the condition for each one. We use two nested loops to define the start (i) and end (j) of a subarray. For a fixed starting point i, we expand the subarray by moving j from i to the end of the array. To efficiently count distinct elements for the expanding subarray nums[i...j], we use a HashSet. As we add nums[j] to the set, we check its size. If the size becomes exactly k, we've found a 'good' subarray and increment our counter. This process is repeated for all possible starting positions i.

import java.util.HashSet;import java.util.Set; class Solution {    public int subarraysWithKDistinct(int[] nums, int k) {        int count = 0;        int n = nums.length;        for (int i = 0; i < n; i++) {            Set<Integer> distinctElements = new HashSet<>();            for (int j = i; j < n; j++) {                distinctElements.add(nums[j]);                if (distinctElements.size() == k) {                    count++;                } else if (distinctElements.size() > k) {                    // Optimization: if size exceeds k, no need to check further for this i                    break;                }            }        }        return count;    }}

Complexity

Time

O(N^2). There are two nested loops. The outer loop runs `N` times, and the inner loop runs up to `N` times. The operations inside the inner loop (adding to a `HashSet` and checking its size) take constant time on average.

Space

O(k). The `HashSet` used to store distinct elements for a subarray will hold at most `k+1` elements before the inner loop breaks. In the worst case, `k` can be up to `N`, making it O(N).

Trade-offs

Pros

  • Simple to understand and implement.

  • It is a direct translation of the problem statement into code.

Cons

  • Highly inefficient for larger input sizes, likely resulting in a 'Time Limit Exceeded' (TLE) error on competitive programming platforms.

  • Performs a lot of redundant work by re-calculating the set of distinct elements for overlapping subarrays.

Solutions

class Solution {public  int subarraysWithKDistinct(int[] nums, int k) {    int[] left = f(nums, k);    int[] right = f(nums, k - 1);    int ans = 0;    for (int i = 0; i < nums.length; ++i) {      ans += right[i] - left[i];    }    return ans;  }private  int[] f(int[] nums, int k) {    int n = nums.length;    int[] cnt = new int[n + 1];    int[] pos = new int[n];    int s = 0;    for (int i = 0, j = 0; i < n; ++i) {      if (++cnt[nums[i]] == 1) {        ++s;      }      for (; s > k; ++j) {        if (--cnt[nums[j]] == 0) {          --s;        }      }      pos[i] = j;    }    return pos;  }}

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next →

Next problem