Problem statement

https://leetcode.com/problems/subarrays-with-k-different-integers/

Solution

The idea is the following: for each ending point let us find the length of the longest window with exactly K different integers and the length of the shortest window with exactly K different integers.

We will use sliding window approach here:

  1. First, we define beg and end pointers equal to 0. vals is counter of values inside window and diff is number of different values in window.
  2. On each step we move end pointer one point to the right and update vals and diff
  3. On the next step we move beg pointer to the right until number of different integers inside window becomes <= K (or <= K-1 in another pass).

  4. For the first pass window(K + 1) we continue to expand our window to the left while we have K or less different symbols inside. Because we stopped one step later, we assign out[end] = end - beg + 1 is the length of longest window ending with end point.

    For the second pass window(K) we continue to expand our window to the left while we have K-1 symbols or less, so here we stop just at the moment when we have K symbols first time, so we need to add 1 later.

    Finally, when we calculate for each end how many windows we have, it is exactly difference between two previously calculated numbers. Imagine, that we have [1, 2, 1, 2, 3] for example. Then for end = 3 we have the smallest window ends with this place is [1, 2] and the longest is [1, 2, 1, 2], so we have in total 3 windows, which can be evaluated as 4 - 1 = 4 - 2 + 1, where 4 is length of the longest window, 2 is length of the shortest window and +1 because between a and b there is b - a + 1 numbers.

Complexity

Both time and space complexity is O(n).

Code

class Solution:
    def subarraysWithKDistinct(self, A, K):
        def window(K):
            beg, end, vals, diff = 0, 0, Counter(), 0
            out = [None]*n
            while end < n:
                vals[A[end]] += 1
                if vals[A[end]] == 1: diff += 1
                
                while diff >= K:
                    if vals[A[beg]] == 1: diff -= 1
                    vals[A[beg]] -= 1
                    beg += 1

                out[end] = end - beg + 1
                end += 1
                 
            return sum(out)
        
        n = len(A)
        return window(K+1) - window(K)