子数组统计问题

2488. 统计中位数为 K 的子数组 (opens new window)

子数组:

奇数情况：<k的数个数 = >k的数个数

等价：左侧< + 右侧< = 左侧> + 右侧>

等价：左侧< - 左侧> = 右侧> -右侧<

等价代换：

左侧小于：在 k 左侧且比 k 小的视作 1； 左侧大于：在 k 左侧且比 k 大的视作 −1； 右侧大于：在 k 右侧且比 k 大的视作 1； 右侧小于：在 k 右侧且比 k 小的视作 −1。

偶数情况：<k的数个数+1 = >k的数个数

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        int pos = 0;
        while(nums[pos]!=k){
            pos++;
        }
        var cnt = new HashMap<Integer,Integer>();
        //子数组长度为1 也是答案
        //val是计数
        cnt.put(0,1);
        for(int i=pos-1,x=0;i>=0;i--){
            x += nums[i]<k?1:-1;
            //等价于:cnt.put(x,cnt.getOrDefault(x,0)+1);
            cnt.merge(x,1,Integer::sum);
        }
        //记录cnt[c]和cnt[c+1]个数
        int ans = cnt.get(0)+cnt.getOrDefault(-1,0);
        for(int i=pos+1,x=0;i<n;i++){
            x += nums[i]>k?1:-1;
            ans += cnt.getOrDefault(x,0)+cnt.getOrDefault(x-1,0);
        }
        return ans;
    }
}