[LeetCode321周赛] 6248. 统计中位数为 K 的子数组

LeetCode321周赛

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

https://leetcode.cn/problems/count-subarrays-with-median-k/

Solution

• 因为数组是 1~n 不相同的数，因此要求以 k 为中位数的子数组必须要包含 k。

• 以 k 的位置为分割点得出两个子数组，求第一个子数组的后缀，求第二隔子数组的前缀。求出 大于 k 的值的数量 - 小于 k 的值的数量，分别存入对应的哈希表。

• 若想要 k 为中位数，只需要满足以 k 的前后扩展中 大于 k 的值的数量 - 小于 k 的值的数量 == 0 or 1，然后使用乘法原理计算即可

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        Map<Integer, Integer> left = new HashMap<>();
        Map<Integer, Integer> right = new HashMap<>();
        left.put(0, 1);
        right.put(0, 1);

        int idx = 0;
        // 找到 k 位置
        for (int i = 0; i < n; i++) {
            if (nums[i] == k) {
                idx = i;
                break;
            }
        }

        int c = 0;
        for (int i = idx - 1; i >= 0; i--) {
            c += nums[i] > k ? 1 : -1;
            left.put(c, left.getOrDefault(c, 0) + 1);
        }

        c = 0;
        for (int i = idx + 1; i < n; i++) {
            c += nums[i] > k ? 1 : -1;
            right.put(c, right.getOrDefault(c, 0) + 1);
        }

        AtomicInteger ans = new AtomicInteger();
        left.forEach((key, val) -> {
            ans.addAndGet(left.get(key) * (right.getOrDefault(-key, 0) + right.getOrDefault(1 - key, 0)));
        });

        return ans.get();
    }
}