6248. 统计中位数为 K 的子数组
难度困难15
给你一个长度为 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.length
1 <= n <= 105
1 <= nums[i], k <= n
nums
中的整数互不相同
等价转换(0x3f)
题解:https://leetcode.cn/problems/count-subarrays-with-median-k/solution/deng-jie-zhuan-huan-pythonjavacgo-by-end-5w11/
class Solution:
def countSubarrays(self, nums: List[int], k: int) -> int:
# 展开成一个数学式子
# 中位数 ==> (奇数长度) 小于 k 的数的个数 = 大于 k 的数的个数
# ==> k左侧小于k的个数 + k右侧小于k的个数 = k左侧大于k的个数 + k右侧大于k的个数
# ==> + k左侧小于k的个数 - k左侧大于k的个数 = + k右侧大于k的个数 - k右侧小于k的个数
# 用哈希表将k右边出现的数的次数统计下来,再k的左侧遍历 从右到左寻找个数
# 偶数长度怎么办? 小于+1 = 大于
# 左侧小于 + 右侧小于+1 = 左侧大于 + 右侧大于
# 左侧小于 - 左侧大于+1 = + 右侧大于 - 右侧小于
pos = nums.index(k) # 找到k的位置
cnt = Counter() # int : int 的哈希表
cnt[0] = 1 # 长度为1的个数有一个
c = 0
for i in range(pos+1, len(nums)):
c += 1 if nums[i] > k else -1
cnt[c] += 1
c = 0
ans = cnt[0] + cnt[1] # 加上 k 和 (k,k+i) 的个数
for i in range(pos-1, -1, -1): # 从k左侧从右往左遍历
c += 1 if nums[i] < k else -1
ans += cnt[c] + cnt[c+1]
return ans
中心扩散+哈希表(你好A)
可以注意到条件里说了:nums中的整数互不相同。
也就是说数组里只有一个k,那么我们可以想到,以k为中心向左右两边扩展求子数组。
- 先找到k在nums里的位置idx;
- 从idx开始向左移动,记录路上大于k和小于k的数量,我这里采用正负性来计算,小于k就-1,大于k就+1,并用哈希表left记录下来,初始化left[0]=1;
- 再从idx开始向右移动,也记录路上大于k和小于k的数量,然后根据当前的值sum,看left中有多少数x能使得:sum+x==0或sum+x==1,因为如果为偶数,则大于k的数要比小于k的数多1才是满足条件的答案。
class Solution {
public int countSubarrays(int[] nums, int k) {
int n = nums.length, idx = -1;
//找到 k 对应的下标idx
for(int i = 0; i < n; i++){
if(nums[i] == k) idx = i;
}
if(idx == -1) return -1;
int cnt = 0,sum = 0;
Map<Integer,Integer> left = new HashMap<>();
left.put(0,1);
// 用哈希表记录下idx左侧 大于小于k 出现的情况
for(int i = idx-1; i >= 0; i--){
if(nums[i] < k) sum--;
else sum++;
left.put(sum,left.getOrDefault(sum,0)+1);
}
cnt += left.get(0);
cnt += left.getOrDefault(1,0); // 加上 k 和 (k,k+i) 的个数
sum = 0;
for(int i = idx+1; i < n; i++){
if(nums[i] < k) sum--;
else sum++;
cnt += left.getOrDefault(-1*sum, 0); //sum+x==0
cnt += left.getOrDefault(1-sum, 0); //sum+x==1
}
return cnt;
}
}