2488. 统计中位数为 K 的子数组
前缀和+哈希表。分别统计每个位置大于k和小于k的数字个数的差值，并以此找符合条件的子数组。
对于一个子数组，我们设 k k k 的位置为 i n d e x index index ，大于 k k k 的数字的数量为 r i g h t _ n u m right\_num right_num ，小于 k k k 的数字的数量为 l e f t _ n u m left\_num left_num，二者的差值 d i s = l e f t _ n u m − r i g h t _ n u m dis=left\_num-right\_num dis=left_num−right_num 。则对于任何两个位置 i i i 和 j j j ，如果满足 i < i n d e x < = j i<index<=j i<index<=j 且 d i s [ i ] = d i s [ j ] 或 d i s [ i ] = d i s [ j ] + 1 dis[i]=dis[j]\;或\;dis[i]=dis[j]+1 dis[i]=dis[j]或dis[i]=dis[j]+1则该子数组一定是满足条件的子数组。
- 对于 d i s [ i ] = d i s [ j ] dis[i]=dis[j] dis[i]=dis[j]，表示 [ i , j ] [i,j] [i,j] 之间的 l e f t _ n u m = r i g h t _ n u m left\_num=right\_num left_num=right_num ，满足中位数的第一个情况
- 对于 d i s [ i ] = d i s [ j ] + 1 dis[i]=dis[j]+1 dis[i]=dis[j]+1，由于表示 d i s = l e f t _ n u m − r i g h t _ n u m dis=left\_num-right\_num dis=left_num−right_num，意味着 [ i , j ] [i,j] [i,j] 之间的 l e f t _ n u m − r i g h t _ n u m = − 1 left\_num-right\_num=-1 left_num−right_num=−1 ，符合中位数的第二种情况。
那具体应该怎么做呢，我们只遍历一次数组，维护两个前缀和 l e f t _ n u m left\_num left_num 和 r i g h t _ n u m right\_num right_num。由于这个子数组一定要包含 k k k ，所以在遇到 k k k 之前，我们用哈希表 m p mp mp 记录每一个 d i s dis dis 出现的次数；当遇到 k k k 之后，我们就开始计算子数组的数量，对于 k k k 之后的每一个位置，子数组数量为 m p [ d i s ] + m p [ d i s + 1 ] mp[dis]+mp[dis+1] mp[dis]+mp[dis+1]。
class Solution {
public:
int countSubarrays(vector<int>& nums, int k) {
int n=nums.size();
int res=0;
int left_num=0;
int right_num=0;
unordered_map<int,int> mp;
mp[0]++;
int flag=false;
for(int i=0;i<n;i++){
if(nums[i]<k)left_num++;
if(nums[i]==k) flag=true;
if(nums[i]>k) right_num++;
int dis=left_num-right_num;
if(!flag){
mp[dis]++;
}
else{
res+=mp[dis];
res+=mp[dis+1];
// cout<<dis<<" "<<mp[dis]<<" "<<mp[dis+1]<<endl;
}
}
return res;
}
};
之前也有遇到一个前缀和求子集的题目，可以看这里感受一下。
菩萨蛮 【清·纳兰性德】
萧萧几叶风兼雨，离人偏识长更苦。欹枕数秋天，蟾蜍下早弦。 夜寒惊被薄，泪与灯花落。无处不伤心，轻尘在玉琴。
- 长更：长夜。
- 蟾蜍：代指月亮。