Leetcode: 992. Subarrays with K Different Integers

Is there an alternative problem easier to solve? Compute \(c(x)\) which is the number of arrays with at most \(x\) distinct elements. The answer for the original problem is \(c(k) - c(k-1)\). Time and space complexity is \(O(N)\).

from typing import List


class Solution:
    def subarraysWithKDistinct(self, nums: List[int], k: int) -> int:
        N = len(nums)

        def calc(k):
            c = [0] * (N + 1)
            j = d = ans = 0
            for i in range(N):
                c[nums[i]] += 1
                d += 1 if c[nums[i]] == 1 else 0
                while d > k:
                    c[nums[j]] -= 1
                    d -= 1 if c[nums[j]] == 0 else 0
                    j += 1
                ans += i - j + 1
            return ans

        return calc(k) - calc(k - 1)


assert Solution().subarraysWithKDistinct([1, 2, 1, 2, 3], 2) == 7
assert Solution().subarraysWithKDistinct([1, 2, 1, 3, 4], 3) == 3