class Solution {
public:
int countSubarrays(vector<int>& nums, int k) {
}
};
2488. 统计中位数为 K 的子数组
给你一个长度为 n 的数组 nums ，该数组由从 1 到 n 的 不同 整数组成。另给你一个正整数 k 。
统计并返回 num 中的 中位数 等于 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 的子数组。
提示：
n == nums.length1 <= n <= 1051 <= nums[i], k <= nnums 中的整数互不相同原站题解
python3 解法, 执行用时: 1520 ms, 内存消耗: 27.1 MB, 提交时间: 2023-02-26 11:17:44
from sortedcontainers import SortedList
class Solution:
def countSubarrays(self, nums: List[int], k: int) -> int:
def cal(mid: int) -> int:
"""中位数 >=mid 的子数组个数"""
curNums = [1 if num >= mid else -1 for num in nums]
res, curSum, sl = 0, 0, SortedList([0])
for num in curNums:
curSum += num
res += sl.bisect_left(curSum)
sl.add(curSum)
return res
return cal(k) - cal(k + 1)
golang 解法, 执行用时: 64 ms, 内存消耗: 7.9 MB, 提交时间: 2023-02-26 11:17:13
func countSubarrays(nums []int, k int) int {
pos := 0
for nums[pos] != k {
pos++
}
cnt, c := map[int]int{0: 1}, 0 // i=pos 的时候 c 是 0，直接记到 cnt 中，这样下面不是大于就是小于
for _, x := range nums[pos+1:] {
if x > k {
c++
} else {
c--
}
cnt[c]++
}
ans := cnt[0] + cnt[1] // i=pos 的时候 c 是 0，直接加到答案中，这样下面不是大于就是小于
for i, c := pos-1, 0; i >= 0; i-- {
if nums[i] < k {
c++
} else {
c--
}
ans += cnt[c] + cnt[c+1]
}
return ans
}
cpp 解法, 执行用时: 88 ms, 内存消耗: 51.4 MB, 提交时间: 2023-02-26 11:17:00
class Solution {
public:
int countSubarrays(vector<int> &nums, int k) {
int pos = find(nums.begin(), nums.end(), k) - nums.begin(), n = nums.size();
unordered_map<int, int> cnt;
cnt[0] = 1; // i=pos 的时候 c 是 0，直接记到 cnt 中，这样下面不是大于就是小于
for (int i = pos + 1, c = 0; i < n; ++i) {
c += nums[i] > k ? 1 : -1;
++cnt[c];
}
int ans = cnt[0] + cnt[1]; // i=pos 的时候 c 是 0，直接加到答案中，这样下面不是大于就是小于
for (int i = pos - 1, c = 0; i >= 0; --i) {
c += nums[i] < k ? 1 : -1;
ans += cnt[c] + cnt[c + 1];
}
return ans;
}
};
java 解法, 执行用时: 14 ms, 内存消耗: 55 MB, 提交时间: 2023-02-26 11:16:47
class Solution {
public int countSubarrays(int[] nums, int k) {
int pos = 0, n = nums.length;
while (nums[pos] != k) ++pos;
var cnt = new HashMap<Integer, Integer>();
cnt.put(0, 1); // i=pos 的时候 c 是 0，直接记到 cnt 中，这样下面不是大于就是小于
for (int i = pos + 1, c = 0; i < n; ++i) {
c += nums[i] > k ? 1 : -1;
cnt.put(c, cnt.getOrDefault(c, 0) + 1);
}
int ans = cnt.get(0) + cnt.getOrDefault(1, 0); // i=pos 的时候 c 是 0，直接加到答案中，这样下面不是大于就是小于
for (int i = pos - 1, c = 0; i >= 0; --i) {
c += nums[i] < k ? 1 : -1;
ans += cnt.getOrDefault(c, 0) + cnt.getOrDefault(c + 1, 0);
}
return ans;
}
}
python3 解法, 执行用时: 96 ms, 内存消耗: 26.7 MB, 提交时间: 2023-02-26 11:16:27
class Solution:
def countSubarrays(self, nums: List[int], k: int) -> int:
pos = nums.index(k)
cnt = defaultdict(int)
cnt[0] = 1 # i=pos 的时候 c 是 0，直接记到 cnt 中，这样下面不是大于就是小于
c = 0
for i in range(pos + 1, len(nums)):
c += 1 if nums[i] > k else -1
cnt[c] += 1
ans = cnt[0] + cnt[1] # i=pos 的时候 c 是 0，直接加到答案中，这样下面不是大于就是小于
c = 0
for i in range(pos - 1, -1, -1):
c += 1 if nums[i] < k else -1
ans += cnt[c] + cnt[c + 1]
return ans