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

题目描述

leetcode 困难题

给你一个长度为 n 的数组 nums ，该数组由从 1 到 n 的 不同 整数组成。另给你一个正整数 k 。

统计并返回 num 中的 中位数 等于 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] 。

提示1：

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

等价转换

由于是求中位数等于 $k$ 的子数组的个数，所以可以把数组中大于 $k$ 的元素转换成 $1$ ，小于的转换成 $-1$ ，自身视为 $0$ 。经过转换后利用类似前缀和的方式进行求解。

我们枚举子数组的右端点，由于子数组必须包含 $k$ ，右端点必须大于等于 $kIndex$ ，前缀和左端点必须小于 $kIndex$ 。对于每个右端点，设其前缀和为 $sumA$ ，只需要找到一个合法的左端点，满足前缀和等于 $sumA$ 或者 $sumA - 1$ 即可（后者为偶数长度情况，注意这里由于 $k$ 自身被视为 $0$ ，所以不会出现奇数长度子数组和为 $1$ 的情况）。于是我们顺序遍历原数组，当尚未出现 $k$ 时，视为合法的左端点，将其前缀和保存到哈希表中，出现 $k$ 以后，视为右端点，此时不再将前缀和保存，而是从哈希表中检索对应值的个数并累加到答案中。

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        int[] preSum = new int[n];
        Map<Integer, Integer> cnt = new HashMap<>();
        cnt.put(0, 1);
        int kIndex = Integer.MAX_VALUE;
        int ans = 0;
        for(int i = 0; i < n; i++){
            int sign = sign(nums[i], k);
            preSum[i] = i == 0 ? sign : preSum[i - 1] + sign;
            kIndex = nums[i] == k ? i : kIndex;
            if(i < kIndex){
                cnt.merge(preSum[i], 1, Integer::sum);
                continue;
            }
            ans += cnt.getOrDefault(preSum[i], 0);
            ans += cnt.getOrDefault(preSum[i] - 1, 0);
        }
        return ans;
    }

    private int sign(int num, int k){
        if(num == k){
            return 0;
        }
        return num > k ? 1 : -1;
    }
}