给你一个长度为 n 的数组 nums ,该数组由从 1 到 n 的 不同 整数组成。另给你一个正整数 k 。
统计并返回 nums 中的 中位数 等于 k 的非空子数组的数目。
注意:
数组的中位数是按 递增 顺序排列后位于 中间 的那个元素,如果数组长度为偶数,则中位数是位于中间靠 左 的那个元素。
例如,[2,3,1,4] 的中位数是 2 ,[8,4,3,5,1] 的中位数是 4 。
子数组是数组中的一个连续部分。
示例 1:
输入:nums = [3,2,1,4,5], k = 4
输出:3
解释:中位数等于 4 的子数组有:[4]、[4,5] 和 [1,4,5] 。
示例 2:
输入:nums = [2,3,1], k = 3
输出:1
解释:[3] 是唯一一个中位数等于 3 的子数组。
前缀+哈希
cpp
class Solution {
public:
int countSubarrays(vector<int>& nums, int k) {
//奇数情况: 小 中 大
//偶数情况, 小 中 大 大
int n = nums.size(),sum = 0,ans = 0;
int index = -1;
for(int i = 0;i < n;i++){
if(nums[i] == k){
index = i;
break;
}
}
unordered_map<int, int> group = {{0,1}};
for(int i = 0;i < n;i++){
if(nums[i] < k) nums[i] = -1;
else if(nums[i] > k) nums[i] = 1;
else nums[i] = 0;
sum += nums[i];
if(i < index){
group[sum]++;
}else{
ans = ans + group[sum] + group[sum - 1];
}
}
return ans;
}
};
时间复杂度:O(n)
空间复杂度:O(n)
cpp
for(int i = 0;i < n;i++){
if(nums[i] == k){
index = i;
break;
}
}
先找到k的位置,方便下面的计算。
cpp
if(nums[i] < k) nums[i] = -1;
else if(nums[i] > k) nums[i] = 1;
else nums[i] = 0;
将大于k的值转换成1,小于k的值转换成-1,然后等于k的值等于0。
cpp
sum += nums[i];
if(i < index){
group[sum]++;
}else{
ans = ans + group[sum] + group[sum - 1];
}
代码中的关键逻辑部分,当i在k左边的时候,我们计算前缀和并储存到哈希表中。当i在k右边的时候,在哈希表中查找是否有对应的前缀和,如果找到了,说明存在以k为中位数的子段,而且该子段元素是奇数个。然后再查找是否有sum-1这个前缀和,如果找到了,说明存在以k为中位数的子段,而且该子段元素是偶数个。将他们记录到ans中即可。