Subarray Distinct Values

Given an array of n integers, your task is to calculate the number of subarrays that have at most k distinct values.

Example:

Input: n=5, k=2, arr={1,2,3,1,1}
Output: 10

Approach

C++

#include <bits/stdc++.h>
using namespace std;

long long subarrayDistinctValues(long long n,long long k,
                         vector<long long> &arr) 
{
 
    map<long long,long long> mp;
    long long ans = 0;
    for (long long i = 0, j = 0; j < n; j++) {
        if (mp.size() == k && !mp.count(arr[j]))
         {
            while ((long long)mp.size() == k) {
                mp[arr[i]]--;
                if (mp[arr[i]] == 0)
                    mp.erase(arr[i]);
                i++;
            }
        }
        mp[arr[j]]++;
        ans += j - i + 1;
    }
    return ans ;
}

int main() {
  
    long long n=5, k=2; 
    
    vector<long long> arr={1,2,3,1,1};
    
   cout<<subarrayDistinctValues(n,k,arr)<<"\n";

   return 0;
    
}