来源:力扣(LeetCode)
描述:
给你一个长度为 n
的数组 nums
,该数组由从 1
到 n
的 不同 整数组成。另给你一个正整数 k 。
统计并返回 nums
中的 中位数 等于 k
的非空子数组的数目。
注意:
[2,3,1,4]
的中位数是 2
,[8,4,3,5,1]
的中位数是 4 。示例 1:
输入:nums = [3,2,1,4,5], k = 4
输出:3
解释:中位数等于 4 的子数组有:[4]、[4,5] 和 [1,4,5] 。
示例 2:
输入:nums = [2,3,1], k = 3
输出:1
解释:[3] 是唯一一个中位数等于 3 的子数组。
提示:
方法:前缀和
由于数组 nums 的长度是 n,数组由从 1 到 n 的不同整数组成,因此数组中的元素各不相同,满足 1 ≤ k ≤ n 的正整数 k 在数组中恰好出现一次。
用 kIndex 表示正整数 k 在数组 nums 中的下标。根据中位数的定义,中位数等于 k 的非空子数组应满足以下条件:
为了计算每个子数组中的大于 k 的元素个数与小于 k 的元素个数之差,需要将原始数组做转换,将大于 k 的元素转换成 1,小于 k 的元素转换成 −1,等于 k 的元素转换成 0,转换后的数组中,每个子数组的元素和为对应的原始子数组中的大于 k 的元素个数与小于 k 的元素个数之差。
为了在转换后的数组中寻找符合要求的子数组,可以计算转换后的数组的前缀和,根据前缀和寻找符合要求的子数组。规定空前缀的前缀和是 0 且对应下标 −1。如果存在下标 left 和 right 满足 −1 ≤ left < kIndex ≤ right < n 且下标 right 处的前缀和与下标 left 处的前缀和之差为 0 或 1,则等价于下标范围 [left + 1, right] 包含下标 kIndex 且该下标范围的转换后的子数组的元素和为 0 或 1,对应该下标范围的原始子数组的中位数等于 k。
根据上述分析,可以从左到右遍历数组 nums,遍历过程中计算转换后的数组的前缀和,并计算中位数等于 k 的非空子数组的数目。
使用哈希表记录转换后的数组的每个前缀和的出现次数。由于空前缀的前缀和是 0,因此首先将前缀和 0 与次数 1 存入哈希表。
用 sum 表示转换后的数组的前缀和,遍历过程中维护 sum 的值。对于 0 ≤ i < n,当遍历到下标 i 时,更新 sum 的值,然后执行如下操作:
遍历结束之后,即可得到中位数等于 k 的非空子数组的数目。
代码:
class Solution {
public:inline int sign(int num) {if (num == 0) {return 0;}return num > 0 ? 1 : -1;}int countSubarrays(vector& nums, int k) {int n = nums.size();int kIndex = -1;for (int i = 0; i < n; i++) {if (nums[i] == k) {kIndex = i;break;}}int ans = 0;unordered_map counts;counts[0] = 1;int sum = 0;for (int i = 0; i < n; i++) {sum += sign(nums[i] - k);if (i < kIndex) {counts[sum]++;} else {int prev0 = counts[sum];int prev1 = counts[sum - 1];ans += prev0 + prev1;}}return ans;}
};
执行用时:96 ms, 在所有 C++ 提交中击败了56.88%的用户
内存消耗:51.3 MB, 在所有 C++ 提交中击败了78.44%的用户
复杂度分析
时间复杂度:O(n),其中 n 是数组 nums 的长度。遍历数组寻找正整数 k 所在下标需要 O(n) 的时间,遍历数组计算中位数等于 k 的非空子数组数目需要 O(n) 的时间。
空间复杂度:O(n),其中 n 是数组 nums 的长度。哈希表需要 O(n) 的空间。
author:LeetCode-Solution