2488. Count Subarrays With Median K

Posted on Edited on In

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:

1
2
3
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:

1
2
3
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 <= 10^5
  • 1 <= nums[i], k <= n
  • The integers in nums are distinct.

Hints/Notes

Solution

Language: C++

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution {
public:
    int countSubarrays(vector<int>& nums, int k) {
        int pos = find(nums.begin(), nums.end(), k) - nums.begin(), n = nums.size();
        unordered_map<int, int> count{{0, 1}};
        for (int i = pos - 1, x = 0; i >= 0; --i) {
            x += nums[i] < k ? 1 : -1;
            ++count[x];
        }
        int ans = count[0] + count[-1];
        for (int i = pos + 1, x = 0; i < n; i++) {
            x += nums[i] > k ? 1 : -1;
            ans += count[x] + count[x - 1];
        }
        return ans;
    }
};