# Count Subarrays With Median K
**Difficulty:** HARD
[External](https://leetcode.com/problems/count-subarrays-with-median-k)
Canonical: https://scaleengineer.com/dsa/problems/count-subarrays-with-median-k
**Patterns:** [Prefix Sum](https://scaleengineer.com/dsa/patterns/prefix-sum)
**Data structures:** Array, Hash Table
**Companies:** [Confluent](https://scaleengineer.com/companies/confluent), [thoughtspot](https://scaleengineer.com/companies/thoughtspot)
---
## Problem
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 <= 105`
* `1 <= nums[i], k <= n`
* The integers in `nums` are distinct.

# Approaches
## Brute Force Enumeration
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.
**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.
**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.
### Explanation
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:
1.  Extract the elements from `nums[i]` to `nums[j]` into a new, temporary array.
2.  Sort this new array to easily find the median.
3.  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.
4.  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:
```java
import java.util.Arrays;

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        int count = 0;
        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                int len = j - i + 1;
                int[] sub = new int[len];
                // Copy subarray
                System.arraycopy(nums, i, sub, 0, len);
                
                // Sort to find median
                Arrays.sort(sub);
                
                int median;
                if (len % 2 == 1) {
                    median = sub[len / 2];
                } else {
                    median = sub[len / 2 - 1];
                }
                
                if (median == k) {
                    count++;
                }
            }
        }
        return count;
    }
}
```
### 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`.

## Prefix Balance with Hash Map
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.
**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`.
**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.
### Explanation
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:
1.  **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).

2.  **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.

```java
import java.util.HashMap;
import java.util.Map;

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        int kIndex = -1;
        for (int i = 0; i < n; i++) {
            if (nums[i] == k) {
                kIndex = i;
                break;
            }
        }

        // Map: balance -> frequency
        Map<Integer, Integer> counts = new HashMap<>();
        counts.put(0, 1); // For an empty prefix part (when subarray starts at k)

        int balance = 0;
        for (int i = kIndex - 1; i >= 0; i--) {
            if (nums[i] > k) {
                balance++;
            } else {
                balance--;
            }
            counts.put(balance, counts.getOrDefault(balance, 0) + 1);
        }

        int result = 0;
        // Count subarrays ending at or after kIndex
        balance = 0;
        for (int j = kIndex; j < n; j++) {
            if (nums[j] > k) {
                balance++;
            } else if (nums[j] < k) {
                balance--;
            }
            // For a subarray nums[i..j], its balance is balance(i..k-1) + balance(k..j).
            // We need total balance to be 0 or 1.
            // balance(i..k-1) = -balance(k..j) OR balance(i..k-1) = 1 - balance(k..j)
            result += counts.getOrDefault(-balance, 0);
            result += counts.getOrDefault(1 - balance, 0);
        }

        return result;
    }
}
```
### 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`.

# Solutions
### Java

```java
class Solution {
public
  int countSubarrays(int[] nums, int k) {
    int n = nums.length;
    int i = 0;
    for (; nums[i] != k; ++i) {
    }
    int[] cnt = new int[n << 1 | 1];
    int ans = 1;
    int x = 0;
    for (int j = i + 1; j < n; ++j) {
      x += nums[j] > k ? 1 : -1;
      if (x >= 0 && x <= 1) {
        ++ans;
      }
      ++cnt[x + n];
    }
    x = 0;
    for (int j = i - 1; j >= 0; --j) {
      x += nums[j] > k ? 1 : -1;
      if (x >= 0 && x <= 1) {
        ++ans;
      }
      ans += cnt[-x + n] + cnt[-x + 1 + n];
    }
    return ans;
  }
}

```

### CPP

```cpp
class Solution {
public:
  int countSubarrays(vector<int> &nums, int k) {
    int n = nums.size();
    int i = find(nums.begin(), nums.end(), k) - nums.begin();
    int cnt[n << 1 | 1];
    memset(cnt, 0, sizeof(cnt));
    int ans = 1;
    int x = 0;
    for (int j = i + 1; j < n; ++j) {
      x += nums[j] > k ? 1 : -1;
      if (x >= 0 && x <= 1) {
        ++ans;
      }
      ++cnt[x + n];
    }
    x = 0;
    for (int j = i - 1; ~j; --j) {
      x += nums[j] > k ? 1 : -1;
      if (x >= 0 && x <= 1) {
        ++ans;
      }
      ans += cnt[-x + n] + cnt[-x + 1 + n];
    }
    return ans;
  }
};

```

### Python

```python
class Solution:
    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

```