Find the number of subarrays of odd lengths that have a median equal to k

Question

Find the number of subarrays of odd lengths that have a median equal to k.

For example: array = [5,3,1,4,7,7], k=4 then there are 4 odd length subarrays with 4 as their median: [4], [1,4,7], [5,3,1,4,7], [3,1,4,7,7] therefore return 4 as the answer.

Can anyone please help me with this subarray problem, I'm not sure how to get the output.

Answer 1

@adf_hater 's logic is correct ( because median is middle element so smaller elements has to be equal to number of bigger elements). Here is the code using same logic

int sum(int start, int end, vector<int>& v) {
   
    unordered_map<ll, ll> prevSum;
    int res = 0, currSum = 0;
    for (int i = start; i < end; i++) {
        currSum += v[i];
        if (currSum == 0)
            res++;
        if (prevSum.find(currSum - 0) != prevSum.end())
            res += (prevSum[currSum - 0]);
        prevSum[currSum]++;
    }

    return res  ;
}

void solve(int n, vector<int>&v, int k){
    vector<int>smaller(n, 0), bigger(n, 0), d(n, 0) ;
    k-= 1;
    for(int i = 0 ; i < n; i++)
        smaller[i] = v[i] < v[k];

    for(int i = 0 ; i < n; i++)
        bigger[i] = v[i] > v[k] ;

    for(int i = 0 ; i< n; i++)
    d[i] = smaller[i] - bigger[i] ;
    
    cout<< sum(0, n, d) - sum(0, k, d) - sum(k+1, n, d)  ;
}

Answer 2

Recently encountered this problem in Online Assessment. However 'k' is index and 1 <= k <= n where n is length of array.

We have to find how many subarrays have arr[k] as median, also subarray has to be odd length. This is the only hint we need.

Since subarrays are odd length it can be (arr[k]) or (1 element on left and right), (2 elements on left and right), so on...

We can maintain smaller and bigger arrays of length n and populate them as follows:

 if(arr[i] < arr[k])
   smaller[i] = 1;
 else
   smaller[i] = 0;

for bigger elements than arr[k]:

 if(arr[i] > arr[k])
    bigger[i] = 1;
 else
    bigger[i] = 0;

This helps us to find in range i...j where i <= j, count of smaller and bigger elements with respect to arr[k].

For arr[k] to be median in the range [i,j], The following condition has to hold.

  (smaller[j] - smaller[i - 1]) = (bigger[j] - bigger[i - 1])

In other words difference between a number of smaller and bigger elements in the range [i, j] is 0.

we create new array d of length n, such that

  d[i] = smaller[i] - bigger[i]

Now problem reduces to finding the number of subarrays having a sum of 0;

But not all subarrays having sum 0 are useful to us. We don't care about the subarrays that do not include 'k'. So,

  ans = subarray_sum_zero(1, n, d) -  subarray_sum_zero(1, k - 1, d) - subarray_sum_zero(k + 1, n, d)

subarray_sum_zero function finds the number of subarrays in array d. You can find the subarray sum equals k using the map in linear time complexity.

overall runtime complexity is O(n) and space complexity is O(n). It should be able pass the tests n = 1e5.