CSES Solutions - Subarray Distinct Values

import java.util.HashMap;

public class Main {
    
    static long solve(long[] arr, long N, long K) {
        // left and right pointers to mark the start and end of
        // the sliding window
        long left = 0, right = 0;
        // Variable to count how many different numbers we have
        // in the window
        long distinct_count = 0;
        // Variable to store the final result
        long result = 0;

        // Map to keep track of how many times each number
        // appears in the window
        HashMap<Long, Long> frequency = new HashMap<>();

        // Slide the window till the window till the right
        // pointer does not reach the end of the array
        while (right < N) {
            // Check if the current number is new or if its
            // count is zero
            if (!frequency.containsKey(arr[(int)right]) || frequency.get(arr[(int)right]) == 0)
                distinct_count++;

            // Update the count of the current number
            frequency.put(arr[(int)right], frequency.getOrDefault(arr[(int)right], 0L) + 1);

            // If there are more than K distinct numbers, shrink
            // the window from the left
            while (distinct_count > K) {
                // Decrease the count of the number going out of
                // the window
                frequency.put(arr[(int)left], frequency.get(arr[(int)left]) - 1);
                // If its count becomes zero, then there will be
                // one less distinct number in the window
                if (frequency.get(arr[(int)left]) == 0)
                    distinct_count--;
                // Move the left pointer to the right to shrink
                // the window
                left++;
            }

            // Calculate the number of subarrays that end at the
            // current position
            result += right - left + 1;

            // Move the right edge of the window to the right to
            // expand it
            right++;
        }
        // Return the result
        return result;
    }

    public static void main(String[] args) {
        // Sample Input
        long N = 5, K = 2;
        long[] arr = { 1, 2, 1, 2, 3 };

        System.out.println(solve(arr, N, K));

        // We're done!
    }
}