这道题要求统计所有子数组中,中位数等于 k 的子数组个数。

核心思路

  1. 先找到 k 在数组中的位置 pos
  2. 中位数定义(对于奇数长度):排序后中间的数 = k
  3. 等价转换:对于子数组,比 k 小的数个数 = 比 k 大的数个数(奇数长度)或差为1(偶数长度,但题目默认奇数长度?需确认)

实际上 LeetCode 2488 的定义:子数组长度为奇数时,中位数是中间元素;长度偶数时,中位数定义为中间靠左那个。但要注意,题目统计的是中位数为 k 的子数组,所以 k 必须出现在子数组中,并且整个子数组中:

  • 小于 k 的数字个数 = 大于 k 的数字个数(当子数组长度为奇数)
  • 小于 k 的数字个数 = 大于 k 的数字个数 - 1(当子数组长度为偶数时,k 在左中位)

但更简单的统一方法是:将比 k 大的数记为 +1,比 k 小的数记为 -1,等于 k 的记为 0。那么对于包含 k 的子数组,左部分和 + 右部分和 = 0(奇数长度)或 = 1(偶数长度,k 在左中位)。

实际上,题解标准做法是:

算法

  1. 找到 k 的位置 idx
  2. 从 idx 向左扫描,维护「大于 k 的数量 - 小于 k 的数量」的差值,存入哈希表
  3. 从 idx 向右扫描,用同样的方式计算,并与左边组合

代码实现

public int countSubarrays(int[] nums, int k) {
    int n = nums.length;
    int pos = -1;
    // 找到 k 的位置
    for (int i = 0; i < n; i++) {
        if (nums[i] == k) {
            pos = i;
            break;
        }
    }
    
    // map 存储左边出现的平衡值及次数
    Map<Integer, Integer> leftMap = new HashMap<>();
    leftMap.put(0, 1); // 空子数组(只包含 k 自身)
    
    int balance = 0;
    // 向左遍历
    for (int i = pos - 1; i >= 0; i--) {
        if (nums[i] > k) {
            balance++;
        } else if (nums[i] < k) {
            balance--;
        }
        leftMap.put(balance, leftMap.getOrDefault(balance, 0) + 1);
    }
    
    int result = 0;
    balance = 0;
    // 向右遍历并统计
    for (int i = pos; i < n; i++) {
        if (i > pos) { // 避免重复处理 pos 位置
            if (nums[i] > k) {
                balance++;
            } else if (nums[i] < k) {
                balance--;
            }
        }
        
        // 需要左边 balance' = -balance(奇数长度)
        // 或者 balance' = 1 - balance(偶数长度,k 在左边)
        result += leftMap.getOrDefault(-balance, 0);
        result += leftMap.getOrDefault(1 - balance, 0);
    }
    
    return result;
}

解释

  • balance 表示当前区间内「大于 k 的数个数 - 小于 k 的数个数」
  • 当子数组长度为奇数,平衡应为 0,即 leftBalance + rightBalance = 0
  • 当子数组长度为偶数,k 在左中位,应满足 leftBalance + rightBalance = 1
  • leftMap 存储从 k 往左的每种平衡值的出现次数(包括空区间)
  • 向右扩展时,对每个右边界,找左边界使得上面两式成立

时间复杂度:O(n),空间复杂度:O(n)

示例
nums = [3,2,1,4,5], k = 4

  • pos = 3
  • 向左:平衡值变化:-1, -2, -1,得到 leftMap: {0:1, -1:2, -2:1}
  • 向右遍历:
    • i=3, balance=0: 需要 left 0 和 1,得 1+0=1
    • i=4, balance=1: 需要 left -1 和 0,得 2+1=3
  • 总数=4

更多推荐