统计中位数为 K 的子数组

给你一个长度为n的数组nums ，该数组由从1 到n的不同整数组成。另给你一个正整数k。

统计并返回nums中的 中位数等于k的非空子数组的数目。

注意：

数组的中位数是按递增顺序排列后位于中间的那个元素，如果数组长度为偶数，则中位数是位于中间靠左的那个元素。

例如，[2,3,1,4] 的中位数是 2 ，[8,4,3,5,1] 的中位数是 4 。
子数组是数组中的一个连续部分。

示例 ：

输入：nums = [3,2,1,4,5], k = 4
输出：3
解释：中位数等于 4 的子数组有：[4]、[4,5] 和 [1,4,5] 。

思路：

首先，看到中位数我们应该敏感的想到，如果一个数组的中位数是k，那么其一定具有以下性质之一：

• 数组中比k大的数的个数 等于 数组中比k小的数的个数；
• 数组中比k大的数的个数 比 数组中比k小的数的个数 多 1 个。

如果当前我们已知合法的子数组的终点位置，如何快速知道以当前终点位置为终点的合法子数组的个数呢？用前缀和 + 哈希表就可以解决。

为了计算每个子数组中的大于k的元素个数与小于k的元素个数之差，需要将原始数组做转换，将大于k的元素转换成1，小于的元素转换成−1，等于 k 的元素转换成 0。转换后的数组中，每个子数组的元素和为对应的原始子数组中的大于 k 的元素个数与小于 k 的元素个数之差。

为了在转换后的数组中寻找符合要求的子数组，可以计算转换后的数组的前缀和，根据前缀和寻找符合要求的子数组（即子数组的和为0或者1的情况）。

规定空前缀的前缀和是0且对应下标-1。

设k在原始数组中所在位置的下标为idx_k，如果存在一个中位数为k的子数组[left，right]，那么肯定有0 <= left <= idx_x <= right < n且下标right处的前缀和与下标left - 1处的前缀和的差为0或者1。

具体算法：

设pre(前缀和)为0，设idx_k为k在原数组中的位置，将map[0] = 1，对应选择的子数组为左端点时的情况。
从左到右遍历数组，对于 0 <= i < n ，如果nums[i] < k，pre减一；如果nums[i] > k， pre加一 。对于当前位置i来说，如果：

• i < idx_k，说明位置i只能被作为某个（或某几个）合法子数组的左端点，此时可以将哈希表中map[pre]的值加一，将其记录下来。
• i >= idx_k，说明位置i只能被作为某个（或某几个）合法子数组的右端点，此时可以在哈希表中查找，如果表中存在key为pre（对应子数组的和为0的情况）或者pre - 1（对应子数组的和为1的情况）的键值对，那么说明这个右端点可以和之前的map[pre] + map[pre - 1]个左端点构成合法子数组，即将答案加上map[pre] + map[pre - 1]。

遍历结束之后，即可得到中位数等于 k 的非空子数组的数目。

代码：

class Solution {
public:
    int countSubarrays(vector<int>& nums, int k) {
        unordered_map<int, int> mp;
        int n = nums.size();
        int pre = 0, ans = 0, idx = n;
        for(int i = 0; i < nums.size(); ++ i) {
            if(nums[i] == k) {
                idx = i;
                break;
            }
        }
        mp[0] = 1;
        for(int i = 0; i < n; ++ i) {
            if(nums[i] < k) pre--;
            else if(nums[i] > k) pre ++;
            if(i < idx) {  //不能等的原因和前缀和做差有关
            	//此时能作为左端点
                mp[pre] ++;   
            } 
            else {
             	//此时能作为右端点
                ans += mp[pre] + mp[pre - 1];  
                //这里不用判 键为pre和pre-1的键值对 是否存在的原因是如果不存在，值为0，不影响答案。
            }
        }
        return ans;
    }
};

类似的题目：

2588. 统计美丽子数组数目
待更新……