2488. Count Subarrays With Median K

Problem

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.

Solution (Java)

class Solution {
    public int countSubarrays(int[] nums, int k) {
        Map<Integer, Integer> prefixSumOfBalance = new HashMap<>();
        prefixSumOfBalance.put(0, 1); // Dummy value of 0's frequency is 1.
        int ans = 0, runningBalance = 0;
        boolean found = false;
        for (int num : nums) {
            if (num < k) {
                --runningBalance;
            }else if (num > k) {
                ++runningBalance;
            }else {
                found = true;
            }
            if (found) {
                ans += prefixSumOfBalance.getOrDefault(runningBalance, 0) + prefixSumOfBalance.getOrDefault(runningBalance - 1, 0);
            }else {
             // prefixSumOfBalance.merge(runningBalance, 1, Integer::sum); // Similar to the following statement, but it is shorter.
                prefixSumOfBalance.put(runningBalance, prefixSumOfBalance.getOrDefault(runningBalance, 0) + 1);
            }
        }
        return ans;
    }
}

Explain:

nope.

Complexity:

• Time complexity : O(n).
• Space complexity : O(n).