2488. Count Subarrays With Median K

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

 

Solutions

Solution: Hash Table + Prefix Sum

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

 

JavaScript

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
const countSubarrays = function (nums, k) {
  const n = nums.length;
  const balanceMap = new Map();
  const median = nums.indexOf(k);
  let balance = 0;
  let result = 0;

  for (let index = median; index >= 0; index--) {
    const num = nums[index];

    if (num > k) balance += 1;
    if (num < k) balance -= 1;

    const count = balanceMap.get(balance) ?? 0;

    balanceMap.set(balance, count + 1);
  }

  balance = 0;

  for (let index = median; index < n; index++) {
    const num = nums[index];

    if (num > k) balance += 1;
    if (num < k) balance -= 1;

    const count = balanceMap.get(-balance) ?? 0;
    const evenCount = balanceMap.get(1 - balance) ?? 0;

    result += count + evenCount;
  }

  return result;
};