Count Subarrays With Median K

Problem Description

You are given an array nums of size n consisting of distinct integers from 1 to n and a positive integer k. Return the number of non-empty subarrays in nums that have a median equal to k.

Key Insights

• The median of an array is the middle element after sorting the array in ascending order.
• For an even-length array, the median is the left middle element.
• A subarray is a contiguous part of an array.
• To have a median equal to k, the number of elements less than k must be equal to the number of elements greater than k in the subarray.
• The problem can be transformed into counting prefix sums and maintaining balance between counts of elements less than and greater than k.

Space and Time Complexity

Time Complexity: O(n)
Space Complexity: O(n)

Solution

To solve the problem, we will use a prefix sum approach combined with a hashmap to track the balance of counts of elements less than and greater than k.

1. Initialize a variable to keep track of the count of subarrays whose median is k.
2. Iterate through the array while maintaining a balance counter:
   • Increment the balance when an element is less than k.
   • Decrement the balance when an element is greater than k.
   • When encountering k, check how many times the current balance has been seen in the hashmap.
3. Update the hashmap with the current balance count.
4. The result will be the total count of valid subarrays.

Code Solutions