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

题目链接

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.length$
• $1\le n\le 10^5$
• $1\le nums[i],k\le 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)$

C++代码：

class Solution {
public:
    int countSubarrays(vector<int>& 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<int,int> 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}) , 0
        for i in range(idx - 1, -1, -1):
            sum = sum + 1 if nums[i] > k else sum - 1
            cnt[sum] += 1

        ans = cnt[0] + cnt[1]
        sum = 0
        for i in range(idx + 1,n):
             sum = sum + 1 if nums[i] > k else sum - 1
             ans += cnt[-1*sum]
             ans += cnt[1-sum]
        return ans