LeetCode 2488. Count Subarrays With Median K

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