Scale Engineer

Count Subarrays With Median K

Hard
#2270Time: O(n^3 log n). There are O(n^2) subarrays. For each subarray of length `L`, copying takes O(L) and sorting takes O(L log L). The dominant operation is sorting, and in the worst case, `L` is O(n), leading to a total complexity of O(n^2 * n log n) = O(n^3 log n).Space: O(n), as a temporary array of up to size `n` is needed to store a subarray for sorting.2 companies
Patterns
Prefix Sum
Data structures
ArrayHash Table
Companies
Confluentthoughtspot

Prompt

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 <= 105
  • 1 <= nums[i], k <= n
  • The integers in nums are distinct.

Approaches

2 approaches with complexity analysis and trade-offs.

The brute-force approach is the most straightforward way to solve the problem. It involves systematically checking every single non-empty subarray within the given array nums. For each of these subarrays, we determine its median and check if it matches the target value k. If it does, we increment a counter. This process continues until all possible subarrays have been examined.

Algorithm

  • Initialize a counter count to 0.
  • Use a nested loop to generate all possible non-empty subarrays. The outer loop iterates through the start index i from 0 to n-1, and the inner loop iterates through the end index j from i to n-1.
  • For each subarray nums[i..j]:
    • Create a temporary copy of the subarray.
    • Sort the temporary array in ascending order.
    • Determine the median based on the length of the subarray (len = j - i + 1):
      • If len is odd, the median is the element at index len / 2.
      • If len is even, the median is the element at index len / 2 - 1.
    • If the calculated median is equal to k, increment the count.
  • After iterating through all subarrays, return the final count.

Walkthrough

This method relies on generating all contiguous subarrays and then finding the median for each one. The generation is done using two nested loops, one for the start index i and another for the end index j.

For every subarray defined by (i, j), we perform the following steps:

  1. Extract the elements from nums[i] to nums[j] into a new, temporary array.
  2. Sort this new array to easily find the median.
  3. Calculate the median's index based on the subarray's length. According to the problem, for an even-length array, the left of the two middle elements is the median.
  4. Compare the element at the median index with k. If they are equal, we've found a valid subarray, and we increment our result counter.

Here is a code snippet demonstrating this logic:

import java.util.Arrays; class Solution {    public int countSubarrays(int[] nums, int k) {        int n = nums.length;        int count = 0;        for (int i = 0; i < n; i++) {            for (int j = i; j < n; j++) {                int len = j - i + 1;                int[] sub = new int[len];                // Copy subarray                System.arraycopy(nums, i, sub, 0, len);                                // Sort to find median                Arrays.sort(sub);                                int median;                if (len % 2 == 1) {                    median = sub[len / 2];                } else {                    median = sub[len / 2 - 1];                }                                if (median == k) {                    count++;                }            }        }        return count;    }}

Complexity

Time

O(n^3 log n). There are O(n^2) subarrays. For each subarray of length `L`, copying takes O(L) and sorting takes O(L log L). The dominant operation is sorting, and in the worst case, `L` is O(n), leading to a total complexity of O(n^2 * n log n) = O(n^3 log n).

Space

O(n), as a temporary array of up to size `n` is needed to store a subarray for sorting.

Trade-offs

Pros

  • Simple to understand and implement.

  • Correctly solves the problem for small input sizes.

Cons

  • Extremely inefficient due to the repeated sorting of subarrays.

  • The time complexity of O(n^3 log n) is too slow for the given constraints and will result in a 'Time Limit Exceeded' error.

Solutions

class Solution {public  int countSubarrays(int[] nums, int k) {    int n = nums.length;    int i = 0;    for (; nums[i] != k; ++i) {    }    int[] cnt = new int[n << 1 | 1];    int ans = 1;    int x = 0;    for (int j = i + 1; j < n; ++j) {      x += nums[j] > k ? 1 : -1;      if (x >= 0 && x <= 1) {        ++ans;      }      ++cnt[x + n];    }    x = 0;    for (int j = i - 1; j >= 0; --j) {      x += nums[j] > k ? 1 : -1;      if (x >= 0 && x <= 1) {        ++ans;      }      ans += cnt[-x + n] + cnt[-x + 1 + n];    }    return ans;  }}

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next →

Next problem