2488. Count Subarrays With Median K

Problem Link

Description

---

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.

  <ul>
  	<li>For example, the median of <code>[2,3,1,4]</code> is <code>2</code>, and the median of <code>[8,4,3,5,1]</code> is <code>4</code>.</li>
  </ul>
  </li>
  <li>A subarray is a contiguous part of an array.</li>
  

 

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.

Solution

---

Python3

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        N = len(nums)
        
        kIndex = nums.index(k)
        t = nums[kIndex]
        res = 0
        
        ld = rd = 0
        left, right = defaultdict(int), defaultdict(int)
        left[0] = right[0] = 1
        
        for index in range(kIndex + 1, N):
            if nums[index] > t:
                rd += 1
            else:
                rd -= 1
            
            right[rd] += 1
        
        for index in range(kIndex - 1, -1, -1):
            if nums[index] > t:
                ld += 1
            else:
                ld -= 1
            
            left[ld] += 1
        
        for x in left.keys():
            res += left[x] * right[-x]
            res += left[x] * right[-x + 1]
 
        return res