Count Subarrays Where Max Element Appears at Least K Times
Medium
You are given an integer array nums and a positive integer k.
Return the number of subarrays where the maximum element of nums appears at least k times in that subarray.
A subarray is a contiguous sequence of elements within an array.
For example:
Given nums = [1,3,2,3,3], k = 2, the expected output is 6.
The subarrays that contain the element 3 at least 2 times are: [1,3,2,3], [1,3,2,3,3], [3,2,3], [3,2,3,3], [2,3,3] and [3,3].
Given nums = [1,4,2,1], k = 3, the expected output is 0.
No subarray contains the element 4 at least 3 times. Consider edge cases such as an empty input array, what should happen when k is zero, and what to return if k is greater than the number of elements in nums.
Solution
Brute Force Solution
One straightforward approach is to generate all possible subarrays and, for each subarray, count the occurrences of the maximum element. If the count is at least k, we increment our result. This is simple to implement but not very efficient.
def count_subarrays_brute_force(nums, k):
n = len(nums)
count = 0
for i in range(n):
for j in range(i, n):
sub_array = nums[i:j+1]
max_val = max(sub_array)
max_count = sub_array.count(max_val)
if max_count >= k:
count += 1
return count
Time Complexity: O(n^3), where n is the length of the nums array. We have nested loops to generate all subarrays (O(n^2)), and for each subarray, we find the maximum and count its occurrences (O(n)).
Space Complexity: O(n) in the worst case, to store the subarray.
Optimal Solution: Sliding Window
A more efficient approach involves using a sliding window technique. We iterate through the array, expanding the window as long as the count of the maximum element is less than k. When the count reaches k, we calculate how many subarrays end at the current index and add it to the total count. Then, we shift the window by moving the left pointer.
First, find the maximum element in nums. Then use two pointers to maintain a sliding window. When the count of max element is greater or equal to k, then shrink the window from the left.
def count_subarrays_optimal(nums, k):
n = len(nums)
max_val = max(nums)
count = 0
for i in range(n):
max_count = 0
for j in range(i, n):
if nums[j] == max_val:
max_count += 1
if max_count >= k:
count += n - j
break
return count
Time Complexity: O(n), where n is the length of the nums array. We iterate through the array once, and the inner loop does not visit same element of array multiple times.
Space Complexity: O(1). We use a constant amount of extra space.
Edge Cases
- Empty Array: If the input array is empty, the function should return 0.
- k = 0: If k is zero it should return total number of subarrays.
- k greater than array length: If k is larger than the array length, it should return 0 if the max value does not appear
ktimes.
