统计中位数为 K 的子数组

题目

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 <= 10^5
• 1 <= nums[i], k <= n
• nums 中的整数互不相同

题解

方法一：

思路

首先值得注意的一点是每个数都不同，所以k在数组中只会出现一次。

我们还是使用哈希表统计前缀中的信息。

如果在一个区间里k是中位数，那么这个区间中大于k的数a和小于k的数b存在关系a=b或者a=b+1

所以只需要统计前缀中大于k的数和小于k的数的差即可。因此我们可以做转化将数组中大于k的数视为1，小于k的数视为-1，k则视为0。

设前缀和ps[i]为转化后的数组前i个数的和。
这样当两个前缀和ps[i]-ps[j]=0 or 1时，区间[j+1,i]可能是一个合法区间。
为啥是可能？因为这个k有可能并不在这个区间中。区间必须包含了k这个位置才是合法区间。

代码

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