2488. Count Subarrays With Median K

  • Hardness: \(\color{red}\textsf{Hard}\)
  • Ralated Topics: Array、Hash Table、Prefix Sum
  • \(\color{blue}\textsf{Weekly Contest 321}\)

一、題目

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.

二、分析

  • 注意這一題的 median 不是統計上的 median。
  • median 必定是含有 k 的子序列。
  • 我的作法是，找到 k 的位置定為 pivot，其餘的若小於 k，定為 -1，大於 k，定為 +1，並將其用 presum 的方式記錄下來。
  • 根據題目的定義，median == k 只會發生在當 presum[i:j] == 0 或 1 時。又子序列一定要包含 k，所以我們會增加一個條件就是，i >= pivot，j < pivot。其中我們可以事先將 presum[0:pivot] 放入 Hash Map 以供後續使用。

三、解題

1. Hash Map

  • Time complexity: \(O(n)\)
  • Space complexity: \(O(n)\)
int countSubarrays(vector<int>& nums, int k) {
    vector<int> presum = {0};
    int pivot;
    for (int i = 0; i < nums.size(); i++) {
        int back = presum.back();
        if (nums[i] > k) {
            back += 1;
        } else if (nums[i] < k){
            back -= 1;
        } else {
            pivot = i+1;
        }
        presum.push_back(back);
    }
    int res = 0;
    
    unordered_map<int,int> map;
    for (int i = 0; i < pivot; i++) map[presum[i]]++;
    
    for (int i = pivot; i < presum.size(); i++) {
        res += (map[presum[i]] + map[presum[i]-1]);
    }
    return res;
}

回目錄 Catalog