2488. Count Subarrays With Median K

题目 #

You are given an array nums of size n consisting of distinct integers from 1 to n and a positive integer k.

Return the number of non-empty subarrays in nums that have a median equal to k.

Note:

  • The median of an array is the middle element after sorting the array in ascending order. If the array is of even length, the median is the left middle element.
    • For example, the median of [2,3,1,4] is 2, and the median of [8,4,3,5,1] is 4.
  • A subarray is a contiguous part of an array.

Example 1:

Input: nums = [3,2,1,4,5], k = 4
Output: 3
Explanation: The subarrays that have a median equal to 4 are: [4], [4,5] and [1,4,5].

Example 2:

Input: nums = [2,3,1], k = 3
Output: 1
Explanation: [3] is the only subarray that has a median equal to 3.

Constraints:

  • n == nums.length
  • $1 <= n <= 10^5$
  • 1 <= nums[i], k <= n
  • The integers in nums are distinct.

思路1 统计一边大于或小于等于，然后计算另一边 #

分析 #

  • 此题目中的中位数可以推出来
  • 奇数长度下

$$左侧小于+右侧小于 = 左侧大于+右侧大于$$ $$左侧小于-左侧大于 = 右侧大于-右侧小于$$

  • 偶数长度下

$$左侧小于-左侧大于+1 = 右侧大于-右侧小于$$

  • 从中位数开始，向左就是大于为-1，小于为1，累加之后如果左侧和右侧相等就找到一个子数组
  • 那么步骤就是先统计左边，然后去右边找有没有匹配的
  • 边界情况，中位数既要算到左边也要算到右边，可以认为是以中位数为边界，向右找要算上，中位数本身也要统计向左边有几个0或-1

代码 # 

 1func countSubarrays1(nums []int, k int) int {
 2	// 1. 先找中位数的位置
 3	v := 0
 4	for i := range nums {
 5		if nums[i] == k {
 6			v = i
 7			break
 8		}
 9	}
10
11	// 2. 向左开始进行累加，使用map进行存储
12	recordMap := make(map[int]int)
13	recordMap[0] = 1 // 中位数的位置本身就是0
14	tmp := 0
15	for i := v - 1; i >= 0; i-- {
16		if nums[i] > k {
17			tmp--
18		} else {
19			tmp++
20		}
21		if _, ok := recordMap[tmp]; !ok {
22			recordMap[tmp] = 0
23		}
24		recordMap[tmp]++
25	}
26
27	// 向右开始查找
28	result := 0
29	tmp = 0
30	for i := v; i < len(nums); i++ {
31		if nums[i] > k {
32			tmp++
33		} else if nums[i] < k {
34			tmp--
35		}
36		// 左侧小于-左侧大于 = 右侧大于-右侧小于
37		if count, ok := recordMap[tmp]; ok {
38			result += count
39		}
40		// 左侧小于-左侧大于+1 = 右侧大于-右侧小于
41		if count, ok := recordMap[tmp-1]; ok {
42			result += count
43		}
44	}
45	return result
46}

思路2 简化成一次遍历 #

分析 #

  • 还是上面的思路，最左侧其实等于 $\sum 左侧小于 - \sum 左侧大于$
  • 整体统计数据加上 $\sum 左侧大于 - \sum 左侧小于$，那么最左边的就是0，第二个就是大于加一，小于减一，和右边计算是一致的
  • 只需要考虑边界情况，将中位数计算到两边即可，那么一次遍历就解决了问题
  • 第一位是按照0开始计算的，所以每个i都是在算i+1的结果，所以左边的遍历只需要到k所在位置前一个就好

代码 # 

 1func countSubarrays(nums []int, k int) int {
 2	result := 0
 3	v := len(nums)
 4	tmp := 0
 5	recordMap := make(map[int]int)
 6	recordMap[0] = 1 // 第一位是0
 7	for i := range nums {
 8		if nums[i] > k {
 9			tmp++
10		} else if nums[i] < k {
11			tmp--
12		} else {
13			v = i
14		}
15		// 每个i计算的都是下一位的值，所以不需要到中位数所在的位置
16		if i < v {
17			if _, ok := recordMap[tmp]; !ok {
18				recordMap[tmp] = 0
19			}
20			recordMap[tmp]++
21		} else {
22			if count, ok := recordMap[tmp]; ok {
23				result += count
24			}
25			if count, ok := recordMap[tmp-1]; ok {
26				result += count
27			}
28		}
29	}
30	return result
31}