统计中位数为 K 的子数组

You are given an array nums of size n consisting of distinct integers from 1 to n and a positive integer k.

Return the number of non-empty subarrays in nums that have a median equal to k.

Note:

• The median of an array is the middle element after sorting the array in ascending order. If the array is of even length, the median is the left middle element.
  • For example, the median of [2,3,1,4] is 2, and the median of [8,4,3,5,1] is 4.
• A subarray is a contiguous part of an array.

 

Example 1:

Input: nums = [3,2,1,4,5], k = 4
Output: 3
Explanation: The subarrays that have a median equal to 4 are: [4], [4,5] and [1,4,5].

Example 2:

Input: nums = [2,3,1], k = 3
Output: 1
Explanation: [3] is the only subarray that has a median equal to 3.

 

Constraints:

• n == nums.length
• 1 <= n <= 105
• 1 <= nums[i], k <= n
• The integers in nums are distinct.

/*
 * 思路：
 * 1. 使用Java Stream API处理子数组。
 * 2. 遍历所有可能的子数组，检查它们的中位数是否为k。
 */

import java.util.Arrays;
import java.util.stream.IntStream;

public class Solution {
    public int countSubarrays(int[] nums, int k) {
        return (int) IntStream.range(0, nums.length)
            .flatMap(start -> IntStream.range(start, nums.length)
                .map(end -> {
                    int[] subArray = Arrays.copyOfRange(nums, start, end + 1);
                    Arrays.sort(subArray);
                    return subArray[(end - start) / 2];
                })
                .filter(median -> median == k))
            .count();
    }
}

这个题解使用了一种基于平衡计数的策略来高效地统计子数组。首先，找出中位数 k 在数组中的位置 p。然后，从 p 向左扫描，使用 d 来统计当前位置到 p 的平衡状态（大于 k 的元素增加 d，小于 k 的元素减少 d），并在 count 数组中记录每个 d 值出现的次数。接着，从位置 p+1 向右扫描，更新 d 的值，并利用之前记录在 count 数组中的平衡状态来统计符合条件的子数组数量，这是基于左侧的平衡状态可以与右侧的平衡状态结合，形成以 k 为中位数的子数组。最终，通过对所有可能的平衡状态进行统计，得到以 k 为中位数的所有子数组的数量。

O(n)

O(n)