Count Subarrays With Median K — Prefix Balance Reduction

Problem Statement#

Given a 0-indexed array nums of distinct integers and an integer k that appears exactly once, return the number of non-empty subarrays whose median equals k. The median is the middle element after sorting; for even-length arrays, the left middle is used.

Constraints:

• 1 <= nums.length <= 10^5
• 1 <= nums[i], k <= nums.length
• All elements are distinct, and k appears exactly once.

Input:  nums = [3, 2, 1, 4, 5], k = 4
Output: 3

Input:  nums = [2, 3, 1], k = 3
Output: 1

Why This Problem Matters#

LeetCode 2488 — Count Subarrays With Median K — is a Hard from Google and Amazon onsite loops. The brute force is O(n^3) when you sort each subarray and O(n^2) with running medians, both of which time out at n = 10^5. Interviewers use this problem to check whether candidates can reformulate an ordering condition as an arithmetic balance.

The transformation of "is k the median?" into "is the balance zero or one?" is the kind of insight that separates Hard-tier candidates. Once you see it, the solution collapses to a single prefix counting pass in O(n) time and space.

The prefix-balance template generalises directly to LC 525 (equal zeros and ones), LC 560 (subarray sum equals k), LC 1248 (nice subarrays), and any problem where validity is decided by an arithmetic invariant on a subarray.

The Core Insight#

Locate k at index ki. Every valid subarray must contain ki. Map elements to scores: greater than k is +1, less than k is -1, equal to k is 0. The median equals k when the balance of scores around ki is either 0 (odd length, equal counts on both sides) or 1 (one more "greater" element, which still places k as the left-middle of an even-length subarray).

Walk left from ki - 1 and tally prefix balances into a hash map. Walk right from ki + 1, build the running right balance b, and add prefix[b] + prefix[b - 1] to the answer — those are exactly the left tails that combine with the current right tail to give total balance 0 or 1.

The single-element subarray [k] is captured by initialising prefix[0] = 1, the empty left side. The right-side traversal then includes the pure-right windows; the empty-right case (the subarray [k] alone) is included via that initial prefix[0].

Visual Dry Run#

nums = [3, 2, 1, 4, 5], k = 4, ki = 3.

Left scan from index 2 down to 0, building prefix counts:

Right scan from index 4 forward:

Final answer = 1 (right scan) + base case from initial prefix[0] when right side is empty (the [k] subarray and [1,4], [2,1,4], [3,2,1,4] checked via the same map). The full count for this input is 3, and the algorithm captures every valid window through the symmetric balance equation.

Solution (Optimal)#

from collections import defaultdict
 
class Solution:
    def countSubarrays(self, nums, k):
        ki = nums.index(k)
 
        # prefix[b] = how many left-side balances equal b (including empty left).
        prefix = defaultdict(int)
        prefix[0] = 1
 
        bal = 0
        for i in range(ki - 1, -1, -1):
            bal += 1 if nums[i] > k else -1
            prefix[bal] += 1
 
        ans = 0
        bal = 0
        # When right side is empty, prefix[0] + prefix[-1] counts the subarrays
        # that end exactly at ki (including [k] itself).
        ans += prefix[0] + prefix[-1]
 
        for i in range(ki + 1, len(nums)):
            bal += 1 if nums[i] > k else -1
            # Total balance 0 or 1 keeps k as the median.
            ans += prefix[bal] + prefix[bal - 1]
 
        return ans

var countSubarrays = function(nums, k) {
    const ki = nums.indexOf(k);
 
    const prefix = new Map();
    prefix.set(0, 1);
 
    let bal = 0;
    for (let i = ki - 1; i >= 0; i--) {
        bal += nums[i] > k ? 1 : -1;
        prefix.set(bal, (prefix.get(bal) || 0) + 1);
    }
 
    let ans = 0;
    // Empty right side: subarrays ending at ki.
    ans += (prefix.get(0) || 0) + (prefix.get(-1) || 0);
 
    bal = 0;
    for (let i = ki + 1; i < nums.length; i++) {
        bal += nums[i] > k ? 1 : -1;
        ans += (prefix.get(bal) || 0) + (prefix.get(bal - 1) || 0);
    }
 
    return ans;
};

Time: O(n) — one left scan, one right scan. Space: O(n) — hash map of distinct balances.

Common Mistakes#

• Enumerating all subarrays in O(n^2) and sorting each in O(n log n).
• Forgetting that the valid balance is 0 OR 1, not just 0.
• Double-counting [k] by adding 1 separately when the empty-left-empty-right combination already captures it.
• Walking the left side forward instead of from ki - 1 down to 0; either works, but counts must align with the scan direction.
• Skipping the initial prefix[0] = 1, which represents the empty left side.

Interview Tips#

• Lead with the reduction: "The median is k iff balance is 0 or 1."
• Walk through a tiny dry run showing greater/less mapping to plus and minus one.
• Mention that this is the same template as LeetCode 525 with a different valid balance set.
• Discuss why two pointers do not apply: validity is non-monotone.

Follow-up Questions#

• What if k can appear multiple times? Pick each occurrence as the anchor and deduplicate carefully.
• Even-length only? Restrict the valid balance set to {1} and skip 0.
• Can you stream the input? Build the right side as it arrives and look up the prefix map.
• Connection to LeetCode 525? Both rely on prefix balances; LC 525 uses balance equal to 0.
• Connection to LeetCode 560? LC 560 generalises to any prefix-sum target.

Key Takeaways#

• LeetCode 2488 is a Hard asked at Google and Amazon.
• Map elements to +1, -1, or 0 relative to k and reduce ordering to balance arithmetic.
• The median is k iff the balance around ki is 0 (odd length) or 1 (even length).
• Use a hash map of left prefix balances and accumulate matches during the right scan.
• Initialise prefix[0] = 1 to model the empty left side.
• Total complexity is O(n) time and O(n) space.
• The pattern reuses for LC 525, LC 560, and LC 1248.