<p>You are given an array <code>nums</code> of size <code>n</code> consisting of <strong>distinct </strong>integers from <code>1</code> to <code>n</code> and a positive integer <code>k</code>.</p>

<p>Return <em>the number of non-empty subarrays in </em><code>nums</code><em> that have a <strong>median</strong> equal to </em><code>k</code>.</p>

<p><strong>Note</strong>:</p>

<ul>
	<li>The median of an array is the <strong>middle </strong>element after sorting the array in <strong>ascending </strong>order. If the array is of even length, the median is the <strong>left </strong>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>
</ul>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> nums = [3,2,1,4,5], k = 4
<strong>Output:</strong> 3
<strong>Explanation:</strong> The subarrays that have a median equal to 4 are: [4], [4,5] and [1,4,5].
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> nums = [2,3,1], k = 3
<strong>Output:</strong> 1
<strong>Explanation:</strong> [3] is the only subarray that has a median equal to 3.
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>n == nums.length</code></li>
	<li><code>1 &lt;= n &lt;= 10<sup>5</sup></code></li>
	<li><code>1 &lt;= nums[i], k &lt;= n</code></li>
	<li>The integers in <code>nums</code> are distinct.</li>
</ul>

from collections import defaultdict

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        # 1. Find the index of k, and convert array into -1, 0, 1
        idxK = -1
        for i in range(len(nums)):
            if nums[i] < k:
                nums[i] = -1
            elif nums[i] > k:
                nums[i] = 1
            else:
                idxK = i
                nums[i] = 0
        
        # 2. Calculate the prefix sum
        prefix_sum = [0] * (len(nums) + 1)
        for i in range(1, len(prefix_sum)):
            prefix_sum[i] = prefix_sum[i - 1] + nums[i - 1]
        
        # 3. Calculate the three ranges: total [0, n - 1], left [0, idxK - 1], right [idxK + 1, n - 1]
        left = self.helper(0, idxK, prefix_sum)
        right = self.helper(idxK + 1, len(nums), prefix_sum)
        total = self.helper(0, len(nums), prefix_sum)
        
        return total - left - right

    # Helper function to calculate subarray sum = 0 and 1 in a given interval
    def helper(self, l: int, r: int, prefix_sum: List[int]) -> int:
        counter = defaultdict(int)
        ans = 0
        
        for i in range(l, r + 1):
            curr_sum = prefix_sum[i]
            
            # Add to the answer if there are subarrays that have the same sum
            ans += counter[curr_sum]
            
            # Also check for subarrays with sum equal to current sum minus 1
            ans += counter[curr_sum - 1]
            
            # Update the counter for the current sum
            counter[curr_sum] += 1
        
        return ans

27300