Count Subarrays With Median K - Solution & Explanation
Problem Statement
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.
Note:
- The median of an array is the middle element after sorting the array in ascending order. If the array is of even length, the median is the left middle element.
- For example, the median of
[2,3,1,4]is2, and the median of[8,4,3,5,1]is4.
- For example, the median of
- A subarray is a contiguous part of an array.
Example 1:
Input: nums = [3,2,1,4,5], k = 4 Output: 3 Explanation: The subarrays that have a median equal to 4 are: [4], [4,5] and [1,4,5].
Example 2:
Input: nums = [2,3,1], k = 3 Output: 1 Explanation: [3] is the only subarray that has a median equal to 3.
Constraints:
n == nums.length1 <= n <= 1051 <= nums[i], k <= n- The integers in
numsare distinct.
Approach Overview
Problem Overview: Given an integer array nums and a value k, count the number of subarrays where the median is exactly k. A valid subarray must contain k, and the number of elements greater than k and less than k must satisfy the balance required for k to be the median.
Approach 1: Two Pointers with Median Validation (O(n²) time, O(1) space)
This direct approach enumerates every subarray that contains k. First locate the index of k, then expand subarrays using two pointers while checking whether the median remains k. For each candidate subarray, count elements greater than and less than k, and verify the balance condition that defines a valid median. The implementation relies only on simple iteration over the array, so space usage stays constant. However, repeatedly validating subarrays leads to quadratic time complexity, which becomes slow for large inputs. This method is useful for understanding the median condition and works as a baseline before applying more advanced techniques.
Approach 2: Prefix Sum with Count Balance (O(n) time, O(n) space)
The optimal solution converts the median condition into a balance problem. Treat each element relative to k: assign +1 if a number is greater than k, -1 if smaller, and 0 for k. A subarray has median k when the balance between larger and smaller values is either 0 or 1. Start from the index of k and compute running balances to the right. Store frequency counts of balances in a hash table. Then traverse to the left and look up complementary balances that satisfy the median condition.
This technique leverages a prefix sum-style transformation. Instead of evaluating each subarray independently, you track cumulative differences and reuse previously computed counts. Every lookup is O(1), so the full scan runs in linear time. The key insight is that median constraints can be reduced to a simple equality between prefix balances. This removes the need for repeated subarray checks and scales efficiently even when the array size approaches the upper constraint limits.
Recommended for interviews: Interviewers typically expect the prefix-sum balance solution. The brute-force or two-pointer expansion demonstrates that you understand how median constraints translate into counts of elements greater or smaller than k. The hash-map prefix technique shows stronger algorithmic thinking by converting the problem into a balance lookup and achieving O(n) time.
Approach 1: Prefix Sum with Count Balance
This approach uses a prefix sum-like method to count subarrays effectively without resorting to checking each subarray individually. Instead of focusing directly on the subarrays, we keep track of the number of elements greater and less than k as we iterate, and use a hashmap to maintain balances, which help find the required subarrays efficiently.
This Python function first finds the index of k and initializes a balance counter. As it iterates over the elements after k, it updates the balance according to whether elements are greater than or less than k. The hashmap stores counts of balances, and the result is updated each time a previous balance (or its decrement by one) can be reused. This effectively counts subarrays with the median k without constructing them explicitly.
Complexity
Time Complexity: O(n), as we iterate through the array a constant number of times, and hashmap operations are average O(1).
Space Complexity: O(n), for storing the balance counts in a hashmap.
Approach 2: Two Pointers with Median Validation
This approach uses a two-pointer technique that expands and contracts a window over the array. At each step, it checks if the current window has k as its median using a running sum structure. This direct approach is simpler but can be less optimal than prefix-based methods.
This C implementation uses a nested loop to examine each subarray and checks if k is the median by sorting. This is done by extracting each relevant subarray and using a simple sort-and-compare process. Although this is a direct implementation, it's less efficient due to repeated sorting.
Code
C
Complexity
Time Complexity: O(n^3 log n) due to sorting within the nested loops.
Space Complexity: O(n) for temporary storage of subarrays in each check.
Approach 3: Traversal + Counting
First, we find the position i of the median k in the array, and then start traversing from i to both sides, counting the number of subarrays with a median of k.
Define an answer variable ans, which represents the number of subarrays with a median of k. Initially, ans = 1, which means that there is currently a subarray of length 1 with a median of k. In addition, define a counter cnt, used to count the number of differences between the "number of elements larger than k" and the "number of elements smaller than k" in the currently traversed array.
Next, start traversing to the right from i + 1. We maintain a variable x, which represents the difference between the "number of elements larger than k" and the "number of elements smaller than k" in the current right subarray. If x \in [0, 1], then the median of the current right subarray is k, and the answer variable ans is incremented by 1. Then, we add the value of x to the counter cnt.
Similarly, start traversing to the left from i - 1, also maintaining a variable x, which represents the difference between the "number of elements larger than k" and the "number of elements smaller than k" in the current left subarray. If x \in [0, 1], then the median of the current left subarray is k, and the answer variable ans is incremented by 1. If -x or -x + 1 is also in the counter, it means that there is currently a subarray that spans both sides of i, with a median of k, and the answer variable ans increases the corresponding value in the counter, that is, ans += cnt[-x] + cnt[-x + 1].
Finally, return the answer variable ans.
The time complexity is O(n), and the space complexity is O(n). Where n is the length of the array.
In coding, we can directly open an array of length
2 times n + 1, used to count the difference between the "number of elements larger thank" and the "number of elements smaller thank" in the current array. Each time we add the difference byn, we can convert the range of the difference from[-n, n]to[0, 2n].
Code
Python
Java
C++
Go
TypeScript
Complexity Comparison
| Approach | Complexity |
|---|---|
| Prefix Sum with Count Balance | Time Complexity: Space Complexity: |
| Two Pointers with Median Validation | Time Complexity: Space Complexity: |
| Traversal + Counting | — |
Detailed Complexity Analysis
| Approach | Time | Space | When to Use |
|---|---|---|---|
| Two Pointers with Median Validation | O(n²) | O(1) | Conceptual baseline when learning how median constraints work in subarrays |
| Prefix Sum with Count Balance | O(n) | O(n) | Best general solution for large arrays and typical interview expectations |
Video Solution
2488 Count subarrays with median K l leetcode Weekly Contest 321 • Programming Pathshala • 3,911 views views
Watch 9 more video solutions →Frequently Asked Questions
Is Count Subarrays With Median K easy or hard?
Count Subarrays With Median K Python/Java solution
How to solve Count Subarrays With Median K in O(n)?
What is the best approach for Count Subarrays With Median K?
Is Count Subarrays With Median K asked at Google/Amazon/Meta?
What data structure is used in Count Subarrays With Median K?
What is the time complexity of Count Subarrays With Median K?
Ready to solve this problem?
Practice Count Subarrays With Median K with our built-in code editor and test cases.
Practice on FleetCode