秋高斋

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

发表于 更新于 分类于
本文字数： 899 阅读时长 ≈ 1 分钟
计算前缀和，找到子数组

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
public:
    int countSubarrays(vector<int>& nums, int k) {
        vector<int> sign = vector(nums.size() + 1, 0);
        unordered_map<int,int> sum;
        bool isK = false;
        int result = 0;

        sum[0]++;
        for(int i = 1; i <= nums.size(); i++)
        {
            
            if(nums[i - 1] > k)
            sign[i] = sign[i - 1] + 1;
            else if(nums[i - 1] < k)
            sign[i] = sign[i - 1] + -1;
            else
            {
                sign[i] = sign[i - 1];
                isK = true;
            }
            if(isK)
            {
                result += sum[sign[i]];
                result += sum[sign[i] - 1];
            }
            else
            sum[sign[i]]++;

            //cout << sign[i] << endl;
        }
        return result;
    }
};

要使子数组的中位数为k，那么必须符合下列两个条件之一：

  • 比k大的数和比k小的数，数量一样多。
  • 比k大的数比比k小的数多一个。

那么把数组作如下处理，比k大的数，标注为1；比k小的数，标注为-1；k标记为0。

然后计算前缀和，可以把前缀和认为的“高度”，数组的标记（-1、0、1）是“坡度”。符合条件的子数组，数组的头和数组的尾，“高度”和前面的两个条件对应，只有两种情况：

  • 头尾高度相等。
  • 尾的高度比头高1。

头的索引必然小于等于k的索引，尾的索引必然大于等于k的索引（因为子数组必然包含k）。

把可能的头的“高度”记录在一个哈希表中，每有一个尾，只要找到可能的头，结果就加一。

时间复杂度\(O(n)\)。

空间复杂度\(O(n)\)。