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

1、题干

给你一个长度为 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 <= 105
• 1 <= nums[i], k <= n
• nums 中的整数互不相同

整得有点复杂，错了几次，裂开

2、思路

对于任意包含 k 的子数组符合要求的条件是：假设小于 k 的数有 cl个，大于 k 的数有 cr个，只要满足 cl === cr 或者 cl + 1 === cr 即可

以 k 的下标 ki 为界限，考虑三类子数组情况：

• 左半边子数组符合要求
• 右半边子数组符合要求
• 左右两边子数组组合后符合要求

前两类情况在 $O(n)$ 时间复杂度内可以出结果，第三类情况如果暴力求解，时间复杂度会达到$O(n^2)$，用哈希表优化可以做到 $O(n)$

3、代码

function countSubarrays(nums: number[], k: number): number {
    const ki = nums.indexOf(k), om = new Map(), em = new Map();
    let ans = 0, sum = 0;
    for (let i = ki - 1; i > -1; i--) {
        sum += nums[i] < k ? -1 : 1;
        if (!sum || sum === 1) ans++;

        const map = i % 2 ? om : em;
        map.set(sum, (map.get(sum) || 0) + 1);
    }

    sum = 0;
    for (let i = ki + 1; i < nums.length; i++) {
        sum += nums[i] < k ? -1 : 1;
        if (!sum || sum === 1) ans++;

        let map = i % 2 ? om : em;
        if (map.get(-sum)) ans += map.get(-sum);

        map = i % 2 ? em : om;
        if (map.get(1 - sum)) ans += map.get(1 - sum);
    }

    return ans + 1;
};

4、复杂度

• 时间复杂度：$O(n)$
• 空间复杂度：$O(n)$

5、执行结果