Program to count number of sublists with exactly k unique elements in Python

Sometimes we need to count sublists that contain exactly k unique elements from a given list. This problem can be solved using a sliding window approach with the key insight that sublists with exactly k unique elements equals sublists with "at most k" minus sublists with "at most k-1" unique elements.

So, if the input is like nums = [2, 2, 3, 4] and k = 2, then the output will be 3, as we have the sublists like: [2, 2, 3], [2, 3], [3, 4].

Algorithm Approach

To solve this, we will follow these steps ?

• Create a helper function count() that counts sublists with "at most K" unique elements
• Use a sliding window with two pointers and a frequency map
• For each right pointer position, count all valid sublists ending at that position
• Return count(k) - count(k-1) to get exactly k unique elements

Implementation

from collections import Counter

class Solution:
    def solve(self, nums, k):
        def count(K):
            slot = Counter()
            i = res = 0
            for j, x in enumerate(nums):
                slot[x] += 1
                while len(slot) > K:
                    slot[nums[i]] -= 1
                    if slot[nums[i]] == 0:
                        del slot[nums[i]]
                    i += 1
                res += j - i + 1
            return res
        return count(k) - count(k - 1)

# Test the solution
ob = Solution()
nums = [2, 2, 3, 4]
k = 2
result = ob.solve(nums, k)
print(f"Number of sublists with exactly {k} unique elements: {result}")

Number of sublists with exactly 2 unique elements: 3

How It Works

The algorithm uses the mathematical principle that:

Sublists with exactly k unique elements = Sublists with at most k unique elements - Sublists with at most (k-1) unique elements

The helper function count(K) uses a sliding window technique:

• Maintain a frequency counter for elements in the current window
• Expand the window by moving the right pointer
• When unique elements exceed K, shrink the window from the left
• For each valid window position, add the count of all sublists ending at that position

Example Walkthrough

For nums = [2, 2, 3, 4] and k = 2:

nums = [2, 2, 3, 4]
k = 2

# Let's trace the sublists with exactly 2 unique elements
sublists_with_2_unique = [
    [2, 3],      # indices 0-2
    [2, 2, 3],   # indices 0-2  
    [3, 4]       # indices 2-3
]

print("Valid sublists:")
for sublist in sublists_with_2_unique:
    print(sublist)
print(f"Total count: {len(sublists_with_2_unique)}")

Valid sublists:
[2, 3]
[2, 2, 3]
[3, 4]
Total count: 3

Conclusion

This sliding window approach efficiently counts sublists with exactly k unique elements by using the difference between "at most k" and "at most k-1" counts. The time complexity is O(n) and space complexity is O(k) for the frequency counter.