#2488

Hard

auto_awesomeArray scanning plus hash lookup

LeetCode Problem Workspace

Count Subarrays With Median K

Count the number of subarrays in a given array with median equal to a specified value k.

Problem Statement

You are given an array nums of size n consisting of distinct integers from 1 to n and a positive integer k. Your task is to find how many non-empty subarrays in nums have a median equal to k.

A subarray's median is the middle value when the subarray is sorted. Since the array contains distinct integers, the problem becomes one of identifying subarrays where the position of k allows it to be the median. The challenge lies in efficiently counting such subarrays.

Examples

Example 1

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

Output: 3

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

[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.

Solution Approach

Array Transformation and Cumulative Sum

Convert the array into a new representation where each element is mapped to 1 if it is greater than k, -1 if it is smaller than k, and 0 if it equals k. This transformation helps in simplifying the problem to finding subarrays with a sum of zero using cumulative sum and hash table for fast lookup.

Prefix Sum and Hash Lookup

Use a running prefix sum while iterating through the array. For each index, track how many times this prefix sum has occurred using a hash table. A subarray with a zero median will have a prefix sum that repeats, so efficiently counting the valid subarrays relies on identifying such repeats.

Optimizing with Hash Map

A hash map stores the count of prefix sums. If the current prefix sum matches a previous sum, it indicates that there exists a subarray where the cumulative sum between these indices is zero. This approach drastically reduces the complexity compared to brute force methods.

Complexity Analysis

Metric	Value
Time	Depends on the final approach
Space	Depends on the final approach

The time complexity of the solution depends on the number of elements in the array and the efficiency of the hash map operations. With an optimal approach, the time complexity is O(n), where n is the size of the array. The space complexity also depends on the hash map, resulting in O(n) space for storing the prefix sums.

What Interviewers Usually Probe

The candidate effectively leverages the array transformation to simplify the median problem.
The candidate is comfortable using hash maps to track prefix sums and identify subarrays.
The candidate is able to explain the trade-off between brute force methods and the efficient use of prefix sums and hash lookups.

Common Pitfalls or Variants

Common pitfalls

Not recognizing that the problem can be reduced to finding subarrays with a sum of zero.
Misunderstanding the transformation of the array into 1, -1, and 0 values.
Failing to efficiently count subarrays with hash maps, leading to time complexity issues.

Follow-up variants

Consider adding duplicate values in the array and adjusting the approach accordingly.
Explore the possibility of applying the solution to non-distinct integer arrays.
Modify the problem to allow for non-zero medians in subarrays, which could require different transformation or counting logic.

FAQ

What is the optimal approach for solving 'Count Subarrays With Median K'?

The optimal approach is to transform the array using values 1, -1, and 0, then use a prefix sum and hash map to count subarrays with a sum of zero.

How do you handle large arrays in this problem?

For large arrays, the use of cumulative sum and hash maps ensures an O(n) time complexity, making the solution efficient even for the maximum input size.

What is the time complexity of the 'Count Subarrays With Median K' solution?

The time complexity is O(n) because the algorithm processes each element once and uses efficient hash map lookups for counting subarrays.

What are the key challenges in the 'Count Subarrays With Median K' problem?

The main challenge lies in transforming the array efficiently and using the right data structures, like hash maps, to track prefix sums and count valid subarrays.

How does the prefix sum technique work for this problem?

Prefix sum helps track the cumulative sum of transformed values. When a prefix sum repeats, it indicates the presence of a subarray with a zero sum, which corresponds to a valid subarray with median k.

terminal

Solution

Solution 1: 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$.

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

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

Continue Practicing

#2588 count the number of beautiful subarrays #2615 sum of distances #2251 number of flowers in full bloom