题目链接:2488. 统计中位数为 K 的子数组
方法:前缀和 + 哈希
解题思路
- 根据题意可知,在\(k\)是中位数的子数组中,比\(k\)大的数 \(-\) 比\(k\)小的数 \(=\) \(0\) || \(1\)。那么将两种状态,小于\(k\)置\(-1\),大于\(k\)置\(+1\),计算数组的前缀和\(s\)。
- 由于子数组要包含\(k\),所有左右端点一定在\(k\)值的两边(包括\(k\)),即子数组可以分为两部分\([l, r]\) 和 \([r + 1, i]\),\(i < n\);当两部分的区间和为\(0\) || \(1\)时,即大于\(k\)的值比小于\(k\)的值多\(1\)或相等,此时子数组符合条件。
- 现在我们通过\(hash\)统计右区间\([r + 1, i]\)的区间值的数量(\(hash\)[区间值] = 数量),然后再枚举左边区间,找合适的右区间的数量即可。
代码
class Solution {
public:
int countSubarrays(vector<int>& nums, int k) {
int l = find(nums.begin(), nums.end(), k) - nums.begin() + 1; //初始化子数组的左右端点为数值k的下标 + 1
int r = l, n = nums.size(), ans = 0;
vector<int> s(n + 1);
unordered_map<int, int> Num;
while (r <= n) { // 从左往右, 统计右边区间值
if (nums[r - 1] < k) s[r] = s[r - 1] - 1;
else if(nums[r - 1] > k) s[r] = s[r - 1] + 1;
else s[r] = 0;
Num[s[r]] ++ , r ++;
}
while (l >= 1) { // 从右往左, 枚举左边区间
if (nums[l - 1] < k) s[l] = s[l + 1] - 1;
else if(nums[l - 1] > k) s[l] = s[l + 1] + 1;
else s[l] = 0;
ans += Num[-s[l]] + Num[1 - s[l]];
l -- ;
}
return ans;
}
};
复杂度分析
时间复杂度:\(O(n)\);
空间复杂度:\(O(n)\);