2488. 统计中位数为 K 的子数组-526互联

方法：前缀和 + 哈希

解题思路

根据题意可知，在\(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)\)；