LeetCode 2488. 统计中位数为K的子数组

题目描述

给你一个长度为 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 的子数组。

提示：

• n == nums.length
• 1 <= n <= 105
• 1 <= nums[i], k <= n
• nums 中的整数互不相同

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/count-subarrays-with-median-k

题目分析

看到子数组就想办法套前缀和+哈希表

因为数组中的整数互不相同，所以除了k，其他数字要么比k大，要么比k小。中位数为k的子数组满足性质：

• 如果子数组元素个数为奇数，数组内小于k和大于k的整数数量相等
• 如果子数组元素个数为偶数，数组呢小于k的元素数量比大于k的元素数量少一个

处理题目中的nums数组，记小于k的整数为-1，大于k的整数为1，k为k_（因为子数组中必须包含k），再计算前缀和。此时满足条件的奇数长度数组的元素和为0，偶数长度数组的元素和为1。问题就转化为了求子数组和为0或1的数组个数

class Solution {
public:
    int countSubarrays(vector<int>& nums, int k) {
        int n = nums.size(), k_ = nums.size();
        unordered_map<int, int> mp;
        mp[0] = 1;
        int pre_sum = 0, res = 0;
        for (int i = 0; i < n; ++i) {
            if (nums[i] == k) {
                pre_sum += k_;
            } else if (nums[i] < k) {
                -- pre_sum;
            } else {
                ++ pre_sum;
            }
            if (mp.count(pre_sum - k_)) {
                res += mp[pre_sum - k_];
            } 
            if (mp.count(pre_sum - k_ - 1)) {
                res += mp[pre_sum - k_ - 1];
            }
            mp[pre_sum] ++;
        }
        return res;
    }
};

作者：daydayUppp
链接：https://leetcode.cn/problems/count-subarrays-with-median-k/solution/daydayuppp-zhuan-hua-wen-ti-ha-xi-biao-q-etua/
来源：力扣（LeetCode）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

总结

偷个懒。看到困难级的题就放弃了，学篇题解。自己尝试的时候想把k也当做0计算，发现bug很多，遂放弃。还是不要随便改别人的代码了