BinarySearch 0012 Longest Sublist with K Distinct Numbers

Problem statement

https://binarysearch.com/problems/Longest-Sublist-with-K-Distinct-Numbers/

Solution

Equal to Leetcode 0340. Longest Substring with At Most K Distinct Characters

Complexity

Time complexity is O(n), because in total we have 2n movements of our pointers. Space complexity is O(k).

Code

class Solution:
    def solve(self, k, s):
        window = Counter()
        beg, end, ans, n, dist = 0, 0, 0, len(s), 0
        
        while beg < n and end < n:
            if (dist == k and window[s[end]] > 0) or dist < k:
                if end + 1 < n: window[s[end]] += 1
                if window[s[end]] == 1: dist += 1
                end += 1
                ans = max(ans, end - beg)
            else:
                window[s[beg]] -= 1
                if window[s[beg]] == 0: dist -= 1
                beg += 1
        
        return ans