题目描述
给你一个长度为 n
的数组 nums
,该数组由从 1
到 n
的 不同 整数组成。另给你一个正整数 k
。
统计并返回 num
中的 中位数 等于 k
的非空子数组的数目。
注意:
- 数组的中位数是按 递增 顺序排列后位于 中间 的那个元素,如果数组长度为偶数,则中位数是位于中间靠 左 的那个元素。
- 例如,
[2,3,1,4]
的中位数是2
,[8,4,3,5,1]
的中位数是4
。
- 例如,
- 子数组是数组中的一个连续部分。
样例
输入:nums = [3,2,1,4,5], k = 4
输出:3
解释:中位数等于 4 的子数组有:[4]、[4,5] 和 [1,4,5]。
输入:nums = [2,3,1], k = 3
输出:1
解释:[3] 是唯一一个中位数等于 3 的子数组。
限制
n == nums.length
1 <= n <= 10^5
1 <= nums[i], k <= n
nums
中的整数互不相同。
算法
(部分和,哈希表) $O(n)$
- 由于 $nums$ 中仅有一个 $k$,所以可以根据 $k$ 的位置,利用部分和求解。
- 令大于 $k$ 的数字为 $1$,小于等于 $k$ 的数字为 $-1$。
- 对于每个在 $k$ 或 $k$ 之后的位置 $i$,找到 $k$ 之前的位置 $j$,满足 $[j + 1, i]$ 的部分和为 $0$ 或者 $-1$。
- 可以利用哈希表,在从前往后遍历过程中,记录下来每个位置 $j$ 的前缀和。
时间复杂度
- 遍历数组一次,每个位置仅需要常数的时间记录或者读取哈希表,故时间复杂度为 $O(n)$。
空间复杂度
- 需要 $O(n)$ 的额外空间存储哈希表。
C++ 代码
class Solution {
public:
int countSubarrays(vector<int>& nums, int k) {
const int n = nums.size();
bool found = false;
unordered_map<int, int> h;
h[0] = 1;
int s = 0, ans = 0;
for (int i = 0; i < n; i++) {
s += (nums[i] <= k) ? -1 : 1;
if (nums[i] == k)
found = true;
if (found) ans += h[s] + h[s + 1];
else ++h[s];
}
return ans;
}
};