题目链接

Leetcode.2488 统计中位数为 K 的子数组 Rating ：1999

题目描述

给你一个长度为 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.lengthn == nums.lengthn==nums.length
• 1<=n<=1051 <= n <= 10^51<=n<=105
• 1<=nums[i],k<=n1 <= nums[i], k <= n1<=nums[i],k<=n
• ```nums````中的整数互不相同

解法：前缀和 + 哈希表

等价转换：

• nums[i] > k，把 nums[i]看作 1。
• nums[i] < k，把 nums[i]看作 -1。
• nums[i] == k，把 nums[i]看作 0。

就将求 中位数为k的子数组 转换为求 包括nums[i]==0这一项的且子数组和为 1或者 0的子数组。

所以我们最终的目的就是 求子数组和为 0或者 1的子数组个数（其中必须包含 nums[i]==0这一项）。

我们用 idx(nums[idx] == k)把原数组分为两段 [0 , idx - 1]，[idx + 1 , n - 1]。

我们用一个哈希表 cnt记录子数组和 sum的出现次数，cnt[0] = 1代表 nums[idx]本身也是一个答案。

1. 对于 [0 , idx - 1]，我们从 idx - 1遍历到 0（因为要求的子数组必须包含 nums[idx]），更新子数组和 sum的出现次数。由于只考虑了左边的一段，所以合法的答案有 sum == 0和 sum == 1的子数组个数，即 cnt[0] + cnt[1]。我们先用 ans = cnt[0] + cnt[1]。
2. 对于 [idx + 1 , n - 1]，我们从 idx + 1遍历到 n - 1（因为要求的子数组必须包含 nums[idx]），更新子数组和 sum的出现次数。这时我们需要考虑左右两边，左右的子数组的和分别为 leftSum和rightSum，leftSum + rightSum == 0 or 1。此时我们枚举的是 rightSum，所以只需要让 ans加上 cnt[1 - rightSum]和 cnt[-rightSum]h即可。

时间复杂度： O(n)O(n)O(n)

C++代码：

class Solution {
public:int countSubarrays(vector& nums, int k) {int n = nums.size();int idx = -1;for(int i = 0;i < n;i++){if(nums[i] == k){idx = i;break;}}unordered_map cnt;cnt[0] = 1;int sum = 0;for(int i = idx - 1;i >= 0;i--){if(nums[i] > k) sum++;else sum--;cnt[sum]++;}int ans = cnt[0] + cnt[1];sum = 0;for(int i = idx + 1;i < n;i++){if(nums[i] > k) sum++;else sum--;ans += cnt[-1*sum];ans += cnt[1 - sum];}         return ans;}
};

Python代码：

class Solution:def countSubarrays(self, nums: List[int], k: int) -> int:idx = nums.index(k);n = len(nums)cnt , sum = Counter({0:1}) , 0for i in range(idx - 1, -1, -1):sum = sum + 1 if nums[i] > k else sum - 1cnt[sum] += 1ans = cnt[0] + cnt[1]sum = 0for i in range(idx + 1,n):sum = sum + 1 if nums[i] > k else sum - 1ans += cnt[-1*sum]ans += cnt[1-sum]return ans