Last Updated: 5 May, 2025

Count Good Subarrays

Moderate

Problem statement

Given an integer array 'NUMS' of size 'N' and an integer 'K'.


Return the number of good subarrays of 'NUMS'.


A subarray 'ARR' is good if there are at least 'K' pairs of indices (i, j) such that i < j and ARR[i] == ARR[j].


A subarray is a contiguous non-empty sequence of elements within an array.


For Example :
Let 'N' = 5, 'NUMS' = [1, 2, 1, 1, 3], K = 2.
Consider the subarray [1, 2, 1, 1] (indices 0 to 3 of NUMS). The elements are 1, 2, 1, 1.
Pairs with equal values and i < j within this subarray are:
(1 at index 0, 1 at index 2), (1 at index 0, 1 at index 3), (1 at index 2, 1 at index 3).
There are 3 such pairs. Since 3 >= K (2), this subarray is good.

Consider the subarray [1, 2, 1, 1, 3] (indices 0 to 4 of NUMS). The elements are 1, 2, 1, 1, 3.
Pairs with equal values and i < j within this subarray are:
(1 at index 0, 1 at index 2), (1 at index 0, 1 at index 3), (1 at index 2, 1 at index 3).
There are 3 such pairs. Since 3 >= K (2), this subarray is good.

Consider the subarray [2, 1, 1] (indices 1 to 3 of NUMS). The elements are 2, 1, 1.
Pairs with equal values and i < j within this subarray are:
(1 at index 1, 1 at index 2).
There is 1 such pair. Since 1 < K (2), this subarray is not good.

It can be shown that the only good subarrays are [1, 2, 1, 1] and [1, 2, 1, 1, 3].
Therefore, the total number of good subarrays is 2.
Input Format :
The first line contains two integers, 'N' and 'K'.
The second line contains 'N' integers, the elements of the array 'NUMS'.
Output Format :
Return the number of good subarrays.
Note :
You don't need to print anything. Just implement the given function.
Constraints :
1 <= N <= 10^5
0 <= K <= 10^10
0 <= NUMS[i] <= 10^5

Time Limit: 1 sec

Approaches

01 Approach

Approach:

  • We can iterate through all possible right endpoints of subarrays using a variable 'right'.
  • For each 'right', we need to count how many left endpoints 'left' (where 0 <= left <= right) form a good subarray [left, right].
  • We can maintain a sliding window [left, right] and keep track of the number of pairs of equal elements within this window.
  • As we expand the window by moving 'right', we add the new element NUMS[right]. When adding an element with value 'x', if 'x' already appeared 'c' times in the window [left, right - 1], adding the new 'x' creates 'c' new pairs.
  • We update the count of pairs accordingly.
  • While the number of pairs in the window [left, right] is at least 'K', this window is a good subarray. Any subarray starting at an index 'i' between 'left' and 'right' (inclusive) and ending at 'right' (i.e., [i, right]) will also have at least as many pairs as the subarray [left, right].
  • The number of such starting points 'i' is (right - left + 1). We add this count to our total number of good subarrays.
  • After adding the count, we shrink the window from the left by incrementing 'left' and removing NUMS[left]. When removing an element with value 'x', if 'x' appeared 'c' times in the window [left, right] before removing it, removing one 'x' reduces the number of pairs by (c - 1). We update the pair count accordingly. We continue shrinking from the left until the pair count drops below 'K'.
  • Summing up the counts (right - left + 1) for each 'right' when the window is good gives the total number of good subarrays.

Algorithm:

  • Initialize 'total_good_subarrays' to 0.
  • Initialize 'left' pointer to 0.
  • Initialize 'current_pairs' to 0.
  • Initialize a frequency map 'freq' to store counts of elements in the current window.
  • Iterate 'right' from 0 to N - 1 :
    • Let 'element_at_right' be NUMS[right].
    • Add 'element_at_right' to the window:
      • Add the number of existing elements equal to 'element_at_right' to 'current_pairs'. This is freq[element_at_right] if it exists, otherwise 0.
      • Increment the frequency of 'element_at_right' in the map: freq[element_at_right] = freq.get(element_at_right, 0) + 1.
    • While 'current_pairs' >= K :
      • The window [left, right] is good. All subarrays ending at 'right' and starting from 'left' up to 'right' are good.
      • Add the number of such subarrays (right - left + 1) to 'total_good_subarrays'.
      • Remove 'element_at_left' (NUMS[left]) from the window:
        • Let 'element_at_left' be NUMS[left].
        • Decrement the frequency of 'element_at_left' in the map: freq[element_at_left] -= 1.
        • Subtract the number of pairs involving the removed element from 'current_pairs'. This is freq[element_at_left] (the count after decrementing, which is the number of remaining elements equal to it).
        • Increment 'left' by 1.
  • Return 'total_good_subarrays'.