Problem statement

https://binarysearch.com/problems/Diverse-Words/

Solution

Equal to Leetcode 0992. Subarrays with K Different Integers

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(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)