2488. Count Subarrays With Median K LeetCode Solution

In this guide, you will get 2488. Count Subarrays With Median K LeetCode Solution with the best time and space complexity. The solution to Count Subarrays With Median K problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.

Table of Contents

  1. Problem Statement
  2. Complexity Analysis
  3. Count Subarrays With Median K solution in C++
  4. Count Subarrays With Median K solution in Java
  5. Count Subarrays With Median K solution in Python
  6. Additional Resources
2488. Count Subarrays With Median K LeetCode Solution image

Problem Statement of Count Subarrays With Median K

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.

Complexity Analysis

  • Time Complexity: O(n)
  • Space Complexity: O(n)

2488. Count Subarrays With Median K LeetCode Solution in C++

class Solution {
 public:
  int countSubarrays(vector<int>& nums, int k) {
    const int kIndex = find(nums.begin(), nums.end(), k) - nums.begin();
    int ans = 0;
    unordered_map<int, int> count;

    for (int i = kIndex, balance = 0; i >= 0; --i) {
      if (nums[i] < k)
        --balance;
      else if (nums[i] > k)
        ++balance;
      ++count[balance];
    }

    for (int i = kIndex, balance = 0; i < nums.size(); ++i) {
      if (nums[i] < k)
        --balance;
      else if (nums[i] > k)
        ++balance;
      // The subarray that has balance == 0 or 1 having median equal to k.
      // So, add count[0 - balance] and count[1 - balance] to `ans`.
      ans += count[-balance] + count[1 - balance];
    }

    return ans;
  }
};
/* code provided by PROGIEZ */

2488. Count Subarrays With Median K LeetCode Solution in Java

class Solution {
  public int countSubarrays(int[] nums, int k) {
    final int kIndex = find(nums, k);
    int ans = 0;
    Map<Integer, Integer> count = new HashMap<>();

    for (int i = kIndex, balance = 0; i >= 0; --i) {
      if (nums[i] < k)
        --balance;
      else if (nums[i] > k)
        ++balance;
      count.merge(balance, 1, Integer::sum);
    }

    for (int i = kIndex, balance = 0; i < nums.length; ++i) {
      if (nums[i] < k)
        --balance;
      else if (nums[i] > k)
        ++balance;
      // The subarray that has balance == 0 or 1 having median equal to k.
      // So, add count[0 - balance] and count[1 - balance] to `ans`.
      ans += count.getOrDefault(-balance, 0) + count.getOrDefault(1 - balance, 0);
    }

    return ans;
  }

  private int find(int[] nums, int k) {
    for (int i = 0; i < nums.length; ++i)
      if (nums[i] == k)
        return i;
    throw new IllegalArgumentException();
  }
}
// code provided by PROGIEZ

2488. Count Subarrays With Median K LeetCode Solution in Python

class Solution:
  def countSubarrays(self, nums: list[int], k: int) -> int:
    kIndex = nums.index(k)
    ans = 0
    count = collections.Counter()

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

    balance = 0
    for i in range(kIndex, len(nums)):
      if nums[i] < k:
        balance -= 1
      elif nums[i] > k:
        balance += 1
      # The subarray that has balance == 0 or 1 having median equal to k.
      # So, add count[0 - balance] and count[1 - balance] to `ans`.
      ans += count[-balance] + count[1 - balance]

    return ans
# code by PROGIEZ

Additional Resources

Happy Coding! Keep following PROGIEZ for more updates and solutions.