Count Subarrays With Median K

Count Subarrays With Median 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 \,\text{==}\, nums.length$
• $1 \leq n \leq {10}^{5}$
• $1 \leq nums[i], k \leq n$
• The integers in nums are distinct.

 

解题思路

　　没做出来，菜。

　　关键是要想到下面这个性质，如果一个长度为$n$的数组排序后的中位数是$k$，如果$n$是奇数，那么大于$k$的数的个数应该恰好为$\left\lfloor {\frac{n}{2}} \right\rfloor$，小于$k$的数的个数也恰好为$\left\lfloor {\frac{n}{2}} \right\rfloor$；如果$n$是偶数，那么大于$k$的数的个数应该恰好为$\frac{n}{2}$， 小于$k$的数的个数应该恰好为$\frac{n}{2} - 1$。

　　接着我们可以构造一个数组$a$，如果$\text{nums}[i] > k$，那么$a[i] = 1$；如果$\text{nums}[i] < k$，那么$a[i] = -1$；如果$\text{nums}[i] = k$，那么$a[i] = 0$。任然后对数组$a$求前缀和，得到前缀和数组$s$。如果某个区间$[i, j]$排序后的中位数为$k$，问题就变成了，如果区间长度为奇数那么$s[i] - s[j-1] = 0$；如果区间长度为偶数那么$s[i] - s[j-1] = 1$。

　　因此我们可以枚举区间右端点$i$，根据$i$的奇偶性看看有多少个$j$ ($1 \leq j \leq i$)满足上面的条件，因此可以开个哈希表，每次枚举完$i$后根据$i$的奇偶性来统计$s[i]$（代码是枚举$i$之前统计$s[i-1]$）。

　　这里有个坑要注意的是由于区间的中位数必须要包含$k$这个数，假设$k$在原数组中的下标为$x$，因此在用哈希表统计的时候应该统计必定包含$k$的$s[i]$，也就是说当$i > x$的时候就不能再用哈希表记录$s[i]$了。

　　AC代码如下：

class Solution {
public:
    int countSubarrays(vector<int>& nums, int k) {
        int n = nums.size();
        vector<int> s(n + 1);
        int x;
        for (int i = 1; i <= n; i++) {
            if (nums[i - 1] == k) x = i;
            int t = 0;
            if (nums[i - 1] > k) t = 1;
            else if (nums[i - 1] < k) t = -1;
            s[i] = s[i - 1] + t;
        }
        unordered_map<int, int> mp[2];
        int ret = 0;
        for (int i = 1; i <= n; i++) {
            if (i - 1 < x) mp[i - 1 & 1][s[i - 1]]++;   // 当j - 1 >= x，区间[i, j]就不包含k了
            ret += mp[i & 1][s[i] - 1] + mp[i - 1 & 1][s[i]]; // 区间长度为偶数的情况加上区间长度为奇数的情况
        }
        return ret;
    }
};