Leetcode 2488. Count Subarrays With Median K

[woodyiiiiiii]

这道题我有思路，但我无法用代码十分清晰地表达出来。

仔细细心观察，一开始我想用三种模式(?X?、?XY?、?YX?)来进行判断，再细心化简，发现其实满足要求大于k的数量等于或大于1于小于k的数量即可。

这里就同样分成三个部分了，k左边、k右边、k左边和右边的结合。

然后使用balance来表示大于k和小于k的数值的数量的平衡。同时，这样左右结合的时候用Map来判断，简化时间复杂度，有点类似数组预处理前缀的思路。

关键在于用balance代表两数之和，同时用Map记录，这样存储了关键信息，抛弃了不必要的信息。

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        int idx = -1;

        int ans = 1;
        // find
        for (int i = 0; i < n; ++i) {
            if (nums[i] == k) {
                idx = i;
                break;
            }
        }

        // key means the bal, value means the count
        Map<Integer, Integer> left = new HashMap<>();
        // i - 1 -> 0
        int bal = 0;
        for (int i = idx - 1; i >= 0; --i) {
            bal += nums[i] > k ? 1 : -1;
            left.put(bal, left.getOrDefault(bal, 0) + 1);
            // only consider left part
            if (bal == 0 || bal == 1) {
                ans++;
            }
        }

        // i + 1 -> n - 1
        bal = 0;
        for (int i = idx + 1; i < n; ++i) {
            bal += nums[i] > k ? 1 : -1;
            // only consider right part
            if (bal == 0 || bal == 1) {
                ans++;
            }
            // left part and right part
            if (left.containsKey(-bal)) {
                ans += left.get(-bal);
            }
            if (left.containsKey(-bal + 1)) {
                ans += left.get(-bal + 1);
            }
        }

        return ans;
    }
}

---

• https://leetcode.com/problems/count-subarrays-with-median-k/
• https://leetcode.com/problems/count-subarrays-with-median-k/discuss/2851940/Balance