2488. Count Subarrays With Median K

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.

class Solution {
public:
    int countSubarrays(vector<int>& A, int k) {
        unordered_map<int, int> freq;
        int pos = find(begin(A), end(A), k) - begin(A), res = 0;
        for(int i = pos, now = 0; i < A.size(); i++) {
            if(A[i] < k) now -= 1;
            else if(A[i] > k) now += 1;
            freq[now] += 1;
        }
        for(int i = pos, now = 0; i >= 0; i--) {
            if(A[i] < k) now -= 1;
            else if(A[i] > k) now += 1;
            res += freq[-now] + freq[-now + 1];
        }
        return res;
    }
};