LeetCode 992. Subarrays with K Different Integers

Problem Statement

impl Solution {
    fn compute(a: &Vec<i32>, kk: i32) -> i32 {
        let mut count = vec![0; a.len() + 1];
        let (mut i, mut ans, mut k) = (0, 0, kk);
        for j in 0..a.len() {
            if count[a[j] as usize] == 0 {
                k -= 1;
            }
            count[a[j] as usize] += 1;

            while k < 0 {
                count[a[i] as usize] -= 1;
                i += 1;
                if count[a[i - 1] as usize] == 0 {
                    k += 1;
                }
            }

            ans += j - i  + 1;
        }
        return ans as i32;
    }
    pub fn subarrays_with_k_distinct(a: Vec<i32>, k: i32) -> i32 {
        return Self::compute(&a, k) - Self::compute(&a, k - 1);
    }
}