LeetCode 2488 - Count Subarrays With Median K

LeetCode Problem 2488

Difficulty: 🔴 Hard
Topics: Array, Hash Table, Prefix Sum

Solution

Problem Understanding

This problem asks us to count how many contiguous subarrays have a median equal to a given value k.

The array nums contains every integer from 1 to n exactly once, which means all elements are distinct. That detail is extremely important because it removes ambiguity when comparing values relative to k.

The definition of median used here is slightly unusual for even-length arrays. After sorting the subarray:

• If the length is odd, the median is the middle element.
• If the length is even, the median is the left-middle element.

For example:

• [1,2,3] has median 2
• [1,2,3,4] also has median 2

We need to count every non-empty contiguous subarray whose median is exactly k.

The constraints are large:

• n can be as large as 100000
• A quadratic or cubic solution will not finish in time
• We need something close to linear time

A naive implementation would examine every subarray, sort it, and compute its median. That would be far too slow.

There are several important observations and edge cases:

• Since all numbers are distinct, every number is either less than k, equal to k, or greater than k
• Any valid subarray must contain k
• Single-element subarray [k] is always valid
• Even-length and odd-length subarrays behave slightly differently because of the "left-middle" median definition
• Arrays where k is near the beginning or end can easily cause off-by-one mistakes
• Balance counting must correctly handle both odd and even subarray lengths

The key challenge is transforming the median condition into something we can count efficiently.

Approaches

Brute Force Approach

The brute-force solution is straightforward conceptually.

We generate every possible subarray. For each subarray:

1. Check whether it contains k
2. Sort the subarray
3. Compute the median according to the problem definition
4. Count it if the median equals k

This works because it directly simulates the definition of median.

However, the performance is terrible.

There are O(n^2) subarrays. Sorting each subarray takes up to O(n log n). The total complexity becomes approximately O(n^3 log n) in the worst case.

Even if we optimized median extraction, the solution would still be too slow for n = 100000.

Key Insight for the Optimal Solution

The critical observation is that we do not actually care about the exact ordering of elements. We only care whether elements are:

• Less than k
• Greater than k
• Equal to k

Suppose we transform the array like this:

• Numbers greater than k become +1
• Numbers less than k become -1
• k itself becomes 0

Now consider a subarray containing k.

For k to be the median:

• The number of elements greater than k must equal the number of elements less than k, or
• There can be exactly one extra larger element because of the left-middle rule for even lengths

That means the transformed subarray must have a balance sum of either:

• 0
• 1

This converts the problem into a prefix sum counting problem.

We anchor everything around the position of k, then count how many left and right balance combinations produce valid totals.

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n³ log n)	O(n)	Enumerates and sorts every subarray
Optimal	O(n)	O(n)	Uses balance transformation and prefix counting

Algorithm Walkthrough

Step 1: Find the Position of k

First, locate the index where nums[i] == k.

Every valid subarray must include this index because the median must literally be the value k.

This allows us to center the entire computation around that position.

Step 2: Define the Balance Transformation

We define a running balance:

• Add +1 for values greater than k
• Add -1 for values less than k

We do not include k itself in the balance.

The balance represents:

(# greater than k) - (# less than k)

For a valid median:

• Odd-length subarray requires balance 0
• Even-length subarray requires balance 1

Step 3: Process the Right Side

Starting from the index of k, move rightward.

At each step:

• Update the balance
• Store how many times each balance occurs in a hash map

This map tells us how many right-side segments produce a particular balance.

We initialize the map with:

balance 0 occurs once

This represents taking no elements to the right.

Step 4: Process the Left Side

Now move leftward starting from the position of k.

For each left-side balance, we want the total subarray balance to become either:

• 0
• 1

Suppose current left balance is leftBalance.

Then valid right balances are:

-rightBalance = leftBalance
or
-rightBalance = leftBalance - 1

Equivalently:

rightBalance = -leftBalance
or
rightBalance = 1 - leftBalance

We look these up in the hash map and add their frequencies to the answer.

Step 5: Return the Result

After processing all left-side expansions, the accumulated answer is the total number of valid subarrays.

Why it works

The transformed balance exactly tracks the relationship between elements greater than k and elements less than k.

A subarray has median k precisely when:

• The counts are equal, balance 0
• Or there is one extra greater element, balance 1

By splitting the array around k, we can efficiently count all valid balance combinations using prefix sums and a hash map.

Because every subarray is considered exactly once through its left and right expansions, the algorithm is both correct and linear in time complexity.

Python Solution

from collections import defaultdict
from typing import List

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        n = len(nums)
        k_index = nums.index(k)

        balance_count = defaultdict(int)
        balance_count[0] = 1

        balance = 0

        # Process right side including k
        for i in range(k_index + 1, n):
            if nums[i] > k:
                balance += 1
            else:
                balance -= 1

            balance_count[balance] += 1

        answer = 0
        balance = 0

        # Process left side including k
        for i in range(k_index, -1, -1):
            if nums[i] > k:
                balance += 1
            elif nums[i] < k:
                balance -= 1

            answer += balance_count[-balance]
            answer += balance_count[1 - balance]

        return answer

The implementation begins by locating the index of k. Since every valid subarray must contain k, this index becomes the center of the computation.

The balance_count hash map stores frequencies of balances from the right side. A balance is computed using the transformation:

• +1 for values greater than k
• -1 for values less than k

The first loop processes all suffixes starting immediately after k. Each balance frequency is recorded.

The second loop walks leftward from k. At every step, we compute the current left balance and determine how many right balances can combine with it to form a valid total balance of 0 or 1.

The expression:

balance_count[-balance]

counts subarrays with total balance 0.

The expression:

balance_count[1 - balance]

counts subarrays with total balance 1.

The final answer accumulates all such valid combinations.

Go Solution

package main

func countSubarrays(nums []int, k int) int {
	n := len(nums)

	kIndex := 0
	for i, v := range nums {
		if v == k {
			kIndex = i
			break
		}
	}

	balanceCount := map[int]int{}
	balanceCount[0] = 1

	balance := 0

	// Process right side
	for i := kIndex + 1; i < n; i++ {
		if nums[i] > k {
			balance++
		} else {
			balance--
		}

		balanceCount[balance]++
	}

	answer := 0
	balance = 0

	// Process left side
	for i := kIndex; i >= 0; i-- {
		if nums[i] > k {
			balance++
		} else if nums[i] < k {
			balance--
		}

		answer += balanceCount[-balance]
		answer += balanceCount[1-balance]
	}

	return answer
}

The Go implementation follows exactly the same logic as the Python solution.

The primary difference is that Go uses a native map[int]int instead of Python's defaultdict. Missing keys automatically return the zero value for integers, which simplifies lookup logic.

Since the constraints fit comfortably within 32-bit integer range for balances and counts, standard int types are sufficient.

Slices are used naturally since Go arrays are fixed-size and less convenient for dynamic problems.

Worked Examples

Example 1

Input:

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

Position of k:

k_index = 3

Right Side Processing

Index	Value	Change	Balance	balance_count
initial	-	-	0	{0:1}
4	5	+1	1	{0:1, 1:1}

Left Side Processing

Index	Value	Change	Left Balance	Needed Right Balances	Added Count
3	4	0	0	0, 1	2
2	1	-1	-1	1, 2	1
1	2	-1	-2	2, 3	0
0	3	-1	-3	3, 4	0

Final answer:

3

Valid subarrays:

[4]
[4,5]
[1,4,5]

Example 2

Input:

nums = [2,3,1]
k = 3

Position of k:

k_index = 1

Right Side Processing

Index	Value	Change	Balance	balance_count
initial	-	-	0	{0:1}
2	1	-1	-1	{0:1, -1:1}

Left Side Processing

Index	Value	Change	Left Balance	Needed Right Balances	Added Count
1	3	0	0	0, 1	1
0	2	-1	-1	1, 2	0

Final answer:

1

Valid subarray:

[3]

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Each element is processed at most twice
Space	O(n)	Hash map may store up to n distinct balances

The algorithm performs two linear scans:

• One scan to build right-side balances
• One scan to evaluate left-side combinations

Hash map operations are expected O(1) time, so the total complexity remains linear.

The balance values can vary from -n to +n, meaning the hash map may contain up to O(n) entries.

Test Cases

from typing import List
from collections import defaultdict

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        n = len(nums)
        k_index = nums.index(k)

        balance_count = defaultdict(int)
        balance_count[0] = 1

        balance = 0

        for i in range(k_index + 1, n):
            if nums[i] > k:
                balance += 1
            else:
                balance -= 1

            balance_count[balance] += 1

        answer = 0
        balance = 0

        for i in range(k_index, -1, -1):
            if nums[i] > k:
                balance += 1
            elif nums[i] < k:
                balance -= 1

            answer += balance_count[-balance]
            answer += balance_count[1 - balance]

        return answer

sol = Solution()

assert sol.countSubarrays([3,2,1,4,5], 4) == 3  # provided example
assert sol.countSubarrays([2,3,1], 3) == 1  # provided example

assert sol.countSubarrays([1], 1) == 1  # single element array

assert sol.countSubarrays([2,1], 1) == 1  # k at end
assert sol.countSubarrays([1,2], 1) == 2  # [1], [1,2]

assert sol.countSubarrays([5,4,3,2,1], 3) == 5  # decreasing order
assert sol.countSubarrays([1,2,3,4,5], 3) == 5  # increasing order

assert sol.countSubarrays([4,1,2,3], 4) == 1  # k at beginning
assert sol.countSubarrays([1,2,3,4], 4) == 1  # k at end

assert sol.countSubarrays([2,1,4,3,5], 4) == 4  # mixed ordering

assert sol.countSubarrays([3,1,2], 2) == 1  # only singleton valid

assert sol.countSubarrays([1,3,2,4], 2) == 4  # multiple even-length cases

Test Summary

Test	Why
[3,2,1,4,5], k=4	Validates provided example
[2,3,1], k=3	Small provided example
[1], k=1	Smallest possible input
[2,1], k=1	k at array end
[1,2], k=1	Even-length subarray handling
Increasing array	Tests ordered structure
Decreasing array	Tests reverse ordering
k at beginning	Boundary index handling
k at end	Opposite boundary handling
Mixed ordering	General correctness
Singleton-only valid	Ensures invalid larger subarrays rejected
Multiple even-length cases	Verifies left-middle median rule

Edge Cases

One important edge case occurs when the array contains only one element. In this situation, the single element must equal k, and the only valid subarray is the array itself. Many implementations accidentally mishandle empty balance maps or skip singleton subarrays. This implementation handles it naturally because the balance map starts with {0:1}, and the left-side traversal includes the k index itself.

Another critical edge case is when k appears at the beginning or end of the array. Prefix-sum algorithms frequently introduce off-by-one mistakes in boundary situations. Here, the loops are carefully structured so the right-side processing starts at k_index + 1, while the left-side traversal starts directly at k_index. This guarantees all valid subarrays containing k are considered exactly once.

Even-length subarrays are another subtle source of bugs. Because the problem defines the median as the left-middle element, valid balances are not limited to 0. A balance of 1 is also valid. Forgetting this condition causes many solutions to undercount even-length subarrays such as [4,5] in the first example. The implementation explicitly counts both possibilities using:

balance_count[-balance]
balance_count[1 - balance]

Finally, arrays with highly unbalanced distributions around k can produce many negative or positive balances. Using a hash map instead of an indexed array avoids problems with negative indices and keeps the implementation clean and efficient.