Count Subarrays With Median K

Hard

#2270Time: O(n^3 log n). There are O(n^2) subarrays. For each subarray of length `L`, copying takes O(L) and sorting takes O(L log L). The dominant operation is sorting, and in the worst case, `L` is O(n), leading to a total complexity of O(n^2 * n log n) = O(n^3 log n).Space: O(n), as a temporary array of up to size `n` is needed to store a subarray for sorting.2 companies

Patterns

Prefix Sum

Data structures

Array Hash Table

Companies

Confluent thoughtspot

Prompt

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] is 2, and the median of [8,4,3,5,1] is 4.
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.length
1 <= n <= 10⁵
1 <= nums[i], k <= n
The integers in nums are distinct.

Approaches

2 approaches with complexity analysis and trade-offs.

The brute-force approach is the most straightforward way to solve the problem. It involves systematically checking every single non-empty subarray within the given array nums. For each of these subarrays, we determine its median and check if it matches the target value k. If it does, we increment a counter. This process continues until all possible subarrays have been examined.

Algorithm

Initialize a counter count to 0.
Use a nested loop to generate all possible non-empty subarrays. The outer loop iterates through the start index i from 0 to n-1, and the inner loop iterates through the end index j from i to n-1.
For each subarray nums[i..j]:
- Create a temporary copy of the subarray.
- Sort the temporary array in ascending order.
- Determine the median based on the length of the subarray (len = j - i + 1):
  - If len is odd, the median is the element at index len / 2.
  - If len is even, the median is the element at index len / 2 - 1.
- If the calculated median is equal to k, increment the count.
After iterating through all subarrays, return the final count.

Walkthrough

This method relies on generating all contiguous subarrays and then finding the median for each one. The generation is done using two nested loops, one for the start index i and another for the end index j.

For every subarray defined by (i, j), we perform the following steps:

Extract the elements from nums[i] to nums[j] into a new, temporary array.
Sort this new array to easily find the median.
Calculate the median's index based on the subarray's length. According to the problem, for an even-length array, the left of the two middle elements is the median.
Compare the element at the median index with k. If they are equal, we've found a valid subarray, and we increment our result counter.

Here is a code snippet demonstrating this logic:

1import java.util.Arrays;2 3class Solution {4    public int countSubarrays(int[] nums, int k) {5        int n = nums.length;6        int count = 0;7        for (int i = 0; i < n; i++) {8            for (int j = i; j < n; j++) {9                int len = j - i + 1;10                int[] sub = new int[len];11                // Copy subarray12                System.arraycopy(nums, i, sub, 0, len);13                14                // Sort to find median15                Arrays.sort(sub);16                17                int median;18                if (len % 2 == 1) {19                    median = sub[len / 2];20                } else {21                    median = sub[len / 2 - 1];22                }23                24                if (median == k) {25                    count++;26                }27            }28        }29        return count;30    }31}

Complexity

Time

O(n^3 log n). There are O(n^2) subarrays. For each subarray of length `L`, copying takes O(L) and sorting takes O(L log L). The dominant operation is sorting, and in the worst case, `L` is O(n), leading to a total complexity of O(n^2 * n log n) = O(n^3 log n).

Space

O(n), as a temporary array of up to size `n` is needed to store a subarray for sorting.

Trade-offs

Pros

Simple to understand and implement.
Correctly solves the problem for small input sizes.

Cons

Extremely inefficient due to the repeated sorting of subarrays.
The time complexity of O(n^3 log n) is too slow for the given constraints and will result in a 'Time Limit Exceeded' error.

This efficient approach solves the problem in linear time by reframing it. The central observation is that any subarray with k as its median must contain the element k. We can transform the numbers in the array into +1 (for numbers greater than k), -1 (for numbers less than k), and 0 (for k itself). The median property can then be expressed as a condition on the sum of these transformed values, which we call the 'balance'. By using a hash map to store prefix balances, we can efficiently count the valid subarrays in a single pass over the array.

Algorithm

First, find the index kIndex where nums[kIndex] == k. Any subarray with median k must contain this element.
Define the 'balance' of an element x as +1 if x > k and -1 if x < k. The balance of k is 0.
The median condition translates to: a subarray's total balance must be 0 (for odd length) or 1 (for even length).
Create a HashMap called counts to store the frequencies of balances for subarrays starting to the left of k and ending just before k.
Initialize counts with {0: 1} to represent an empty prefix (the case where a subarray starts at kIndex).
Process the left side: Iterate from i = kIndex - 1 down to 0. Maintain a running balance for the subarray nums[i..kIndex-1]. For each i, update the balance and store its frequency in the counts map.
Count subarrays and process the right side:
- Initialize result by counting subarrays that end at kIndex. A subarray nums[i..kIndex] is valid if its balance is 0 or 1. The balance of nums[i..kIndex] is the same as nums[i..kIndex-1], so we add counts.getOrDefault(0, 0) + counts.getOrDefault(1, 0) to result.
- Iterate from j = kIndex + 1 to n-1. Maintain a running balance_right for the subarray nums[kIndex+1..j].
- For each j, we need to find left parts nums[i..kIndex-1] with balance_left such that balance_left + balance_right is 0 or 1.
- This means balance_left must be -balance_right or 1 - balance_right. We look up these values in our counts map and add the frequencies to result.
Return the final result.

Walkthrough

The core of this method is to split the problem around the index of k, let's call it kIndex. We consider all subarrays nums[i..j] where i <= kIndex <= j.

The balance of such a subarray is the sum of the balance of its left part nums[i..kIndex-1] and its right part nums[kIndex+1..j] (since k itself has a balance of 0).

The algorithm works as follows:

Pre-computation (Left of k): We iterate from kIndex - 1 down to 0. We compute the balance for every prefix nums[i..kIndex-1] and store the frequency of each balance value in a hash map, counts. An initial entry counts[0] = 1 handles subarrays that start at kIndex (i.e., have an empty left part).
Counting (Right of k): We then iterate from kIndex to the end of the array. This pass combines the pre-computed left balances with the right-side balances to count all valid subarrays.
- First, we count subarrays that end at kIndex (nums[i..kIndex]). The balance of such a subarray is simply the balance of its left part nums[i..kIndex-1]. A subarray is valid if its balance is 0 or 1. We can get this count from our counts map.
- Next, for subarrays that extend past kIndex (nums[i..j] where j > kIndex), we calculate the balance of the right part nums[kIndex+1..j]. For each right balance, we find how many left parts have a complementary balance (i.e., balance_left + balance_right is 0 or 1) by looking up counts[-balance_right] and counts[1 - balance_right].

This 'meet-in-the-middle' strategy, centered on k, allows us to avoid re-computation and achieve a linear time solution.

1import java.util.HashMap;2import java.util.Map;3 4class Solution {5    public int countSubarrays(int[] nums, int k) {6        int n = nums.length;7        int kIndex = -1;8        for (int i = 0; i < n; i++) {9            if (nums[i] == k) {10                kIndex = i;11                break;12            }13        }14 15        // Map: balance -> frequency16        Map<Integer, Integer> counts = new HashMap<>();17        counts.put(0, 1); // For an empty prefix part (when subarray starts at k)18 19        int balance = 0;20        for (int i = kIndex - 1; i >= 0; i--) {21            if (nums[i] > k) {22                balance++;23            } else {24                balance--;25            }26            counts.put(balance, counts.getOrDefault(balance, 0) + 1);27        }28 29        int result = 0;30        // Count subarrays ending at or after kIndex31        balance = 0;32        for (int j = kIndex; j < n; j++) {33            if (nums[j] > k) {34                balance++;35            } else if (nums[j] < k) {36                balance--;37            }38            // For a subarray nums[i..j], its balance is balance(i..k-1) + balance(k..j).39            // We need total balance to be 0 or 1.40            // balance(i..k-1) = -balance(k..j) OR balance(i..k-1) = 1 - balance(k..j)41            result += counts.getOrDefault(-balance, 0);42            result += counts.getOrDefault(1 - balance, 0);43        }44 45        return result;46    }47}

Complexity

Time

O(n). Finding `kIndex` takes O(n). The first loop to populate the map takes O(kIndex) time. The second loop for counting takes O(n - kIndex) time. The total time complexity is O(n).

Space

O(n) in the worst case for the hash map, as the balance value can range from `-kIndex` to `kIndex`.

Trade-offs

Pros

Highly efficient with a linear time complexity.
Scales well for large inputs, passing the given constraints.

Cons

The logic is more complex and less intuitive than the brute-force approach.
Requires extra space for the hash map.

Solutions

Solution

1class Solution {2public3  int countSubarrays(int[] nums, int k) {4    int n = nums.length;5    int i = 0;6    for (; nums[i] != k; ++i) {7    }8    int[] cnt = new int[n << 1 | 1];9    int ans = 1;10    int x = 0;11    for (int j = i + 1; j < n; ++j) {12      x += nums[j] > k ? 1 : -1;13      if (x >= 0 && x <= 1) {14        ++ans;15      }16      ++cnt[x + n];17    }18    x = 0;19    for (int j = i - 1; j >= 0; --j) {20      x += nums[j] > k ? 1 : -1;21      if (x >= 0 && x <= 1) {22        ++ans;23      }24      ans += cnt[-x + n] + cnt[-x + 1 + n];25    }26    return ans;27  }28}

1class Solution {2public:3  int countSubarrays(vector<int> &nums, int k) {4    int n = nums.size();5    int i = find(nums.begin(), nums.end(), k) - nums.begin();6    int cnt[n << 1 | 1];7    memset(cnt, 0, sizeof(cnt));8    int ans = 1;9    int x = 0;10    for (int j = i + 1; j < n; ++j) {11      x += nums[j] > k ? 1 : -1;12      if (x >= 0 && x <= 1) {13        ++ans;14      }15      ++cnt[x + n];16    }17    x = 0;18    for (int j = i - 1; ~j; --j) {19      x += nums[j] > k ? 1 : -1;20      if (x >= 0 && x <= 1) {21        ++ans;22      }23      ans += cnt[-x + n] + cnt[-x + 1 + n];24    }25    return ans;26  }27};

1class Solution:2    def countSubarrays(self, nums: List[int], k: int) -> int: i = nums . index(k) cnt = Counter() ans = 1 x = 0 for v in nums[i + 1:]: x += 1 if v > k else - 1 ans += 0 <= x <= 1 cnt[x] += 1 x = 0 for j in range(i - 1, - 1, - 1): x += 1 if nums[j] > k else - 1 ans += 0 <= x <= 1 ans += cnt[- x] + cnt[- x + 1] return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem