力扣(LeetCode)2488. 统计中位数为 K 的子数组(C++)

哈希表

找不到 $O(n)$ 思路？ 试一下等价代换。

数组所有数字大小不同，说明数组中最多有一个 k 。数组的 k 要包含在 子数组 里，为了便于思考，分析奇数长度的子数组：

• 在子数组里，大于 k 的数，和小于 k 的数，二者数量相等时，k 就是中位数。

数量关系写作：左侧大于 + 右侧大于 = 左侧小于 + 右侧小于。 ①
由等式关系，我们想到了前缀和做法。但要开 $4$ 个数组，还可以简化。

简化思路：我们在枚举 k 的左侧元素时，可以维护左侧的所有性质。右侧同理。
数量关系转化为：左侧大于 - 左侧小于 = 右侧小于 - 右侧大于。 ②

k 左侧大于 k 的数，贡献记为 $1$；左侧小于，记为 $-1$ 。
k 右侧小于 k 的数，贡献记为 $1$；右侧大于，记为 $-1$ 。

将这里的贡献，结合前缀和思想。算法流程 ：

1. 从数组的 k 开始往左枚举。遇到大于 k的数，前缀和加 $1$，遇到小于 $k$ 的数，前缀和加 $-1$。在每一次枚举的位置，将前缀和记入哈希表。

2. 从数组的 k 开始往右枚举。按照右侧贡献的规则，维护前缀和。哈希表中相同的前缀和，就对应了一些可以使 k 成为中位数的子数组。

分析偶数长度的子数组：
数量关系： 左侧大于 + 右侧大于 = 左侧小于 + 右侧小于 + 1。 ③
转化为： 左侧大于 - 左侧小于 = 右侧大于 + 右侧小于 + 1。 ④

在算法上，枚举 k 左侧不变；枚举 k 右侧时，找哈希表中等于(右侧前缀和+1)的前缀和。

提示 ： k 的贡献为 $0$ 。

class Solution {
public:
    int countSubarrays(vector<int>& nums, int k) {
        int n = nums.size();
        int pos = find(nums.begin(), nums.end(), k) - nums.begin();
        unordered_map<int, int> mp;
        mp[0] ++;
        for (int i = pos - 1, x = 0; i >= 0; i --) {
            if (nums[i] > k) x ++;
            else x --;
            mp[x] ++;
        }
        int ans = mp[0] + mp[1];
        for (int i = pos + 1, x = 0; i < n; i ++) {
            if (nums[i] < k) x ++;
            else x --;
            ans += mp[x] + mp[x + 1];
        }
        return ans;
    }
};

1. 时间复杂度 : $O(n)$ ， $n$ 是数组的大小，$find$ 函数遍历数组，维护前缀和，二者的时间复杂度是 $O(n)$ 。
2. 空间复杂度 : $O(n)$ ， 哈希表的空间复杂度是 $O(n)$ 。

AC

致语

• 理解思路很重要
• 读者有问题请留言，清墨看到就会回复的。