README.md

2488. 统计中位数为 K 的子数组

English Version

题目描述

给你一个长度为 n 的数组 nums ，该数组由从 1 到 n 的 不同 整数组成。另给你一个正整数 k 。

统计并返回 num 中的 中位数 等于 k 的非空子数组的数目。

注意：

• 数组的中位数是按 递增 顺序排列后位于 中间 的那个元素，如果数组长度为偶数，则中位数是位于中间靠 左 的那个元素。

  <ul>
  	<li>例如，<code>[2,3,1,4]</code> 的中位数是 <code>2</code> ，<code>[8,4,3,5,1]</code> 的中位数是 <code>4</code> 。</li>
  </ul>
  </li>
  <li>子数组是数组中的一个连续部分。</li>
  

 

示例 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 中的整数互不相同

解法

方法一：中心扩展

我们先找到数字 $k$ 所在的位置 $i$，然后从 $i$ 向两边扩展，统计以 $i$ 为中心的子数组个数。

扩展的过程中维护两个变量 $mi$ 和 $mx$，分别表示比 $k$ 小的数字个数和比 $k$ 大的数字个数。

先从 $i+1$ 向右扩展，若出现 $0 \leq mx-mi \leq 1$，此时右侧数组的中位数为 $k$，答案加 $1$。然后我们将 $mx-mi$ 的出现次数存放在数组或哈希表中，以便后续查询。

接着从 $i-1$ 向左扩展，若出现 $0 \leq mx-mi \leq 1$，此时左侧数组的中位数为 $k$，答案加 $1$。然后我们查询数组或哈希表中 $mi-mx$ 以及 $mi-mx+1$ 的出现次数，将其加到答案中。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 为数组长度。

Python3

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        i = next(i for i, v in enumerate(nums) if v == k)
        n = len(nums)
        ans = 1
        d = defaultdict(int)
        mi = mx = 0
        for j in range(i + 1, n):
            if nums[j] < k:
                mi += 1
            else:
                mx += 1
            if 0 <= mx - mi <= 1:
                ans += 1
            d[mx - mi] += 1
        mi = mx = 0
        for j in range(i - 1, -1, -1):
            if nums[j] < k:
                mi += 1
            else:
                mx += 1
            if 0 <= mx - mi <= 1:
                ans += 1
            ans += d[mi - mx] + d[mi - mx + 1]
        return ans

Java

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        int i = 0;
        for (int j = 0; j < n; ++j) {
            if (nums[j] == k) {
                i = j;
                break;
            }
        }
        int ans = 1;
        int[] d = new int[n << 1 | 1];
        int mi = 0, mx = 0;
        for (int j = i + 1; j < n; ++j) {
            if (nums[j] < k) {
                ++mi;
            } else {
                ++mx;
            }
            if (mx - mi >= 0 && mx - mi <= 1) {
                ++ans;
            }
            ++d[mx - mi + n];
        }
        mi = 0;
        mx = 0;
        for (int j = i - 1; j >= 0; --j) {
            if (nums[j] < k) {
                ++mi;
            } else {
                ++mx;
            }
            if (mx - mi >= 0 && mx - mi <= 1) {
                ++ans;
            }
            ans += d[mi - mx + n] + d[mi - mx + 1 + n];
        }
        return ans;
    }
}

C++

class Solution {
public:
    int countSubarrays(vector<int>& nums, int k) {
        int n = nums.size();
        int i = 0;
        for (int j = 0; j < n; ++j) {
            if (nums[j] == k) {
                i = j;
                break;
            }
        }
        int ans = 1;
        int d[n << 1 | 1];
        memset(d, 0, sizeof d);
        int mi = 0, mx = 0;
        for (int j = i + 1; j < n; ++j) {
            if (nums[j] < k) ++mi;
            else ++mx;
            if (mx - mi >= 0 && mx - mi <= 1) ++ans;
            ++d[mx - mi + n];
        }
        mi = 0, mx = 0;
        for (int j = i - 1; ~j; --j) {
            if (nums[j] < k) ++mi;
            else ++mx;
            if (mx - mi >= 0 && mx - mi <= 1) ++ans;
            ans += d[mi - mx + n] + d[mi - mx + n + 1];
        }
        return ans;
    }
};

Go

func countSubarrays(nums []int, k int) int {
	n := len(nums)
	var i int
	for j, v := range nums {
		if v == k {
			i = j
			break
		}
	}
	ans := 1
	d := make([]int, n<<1|1)
	mi, mx := 0, 0
	for j := i + 1; j < n; j++ {
		if nums[j] < k {
			mi++
		} else {
			mx++
		}
		if mx-mi >= 0 && mx-mi <= 1 {
			ans++
		}
		d[mx-mi+n]++
	}
	mi, mx = 0, 0
	for j := i - 1; j >= 0; j-- {
		if nums[j] < k {
			mi++
		} else {
			mx++
		}
		if mx-mi >= 0 && mx-mi <= 1 {
			ans++
		}
		ans += d[mi-mx+n] + d[mi-mx+n+1]
	}
	return ans
}

...