Find the number of sub arrays in the permutation of first N natural numbers such that their median is M

Given an array arr[] containing the permutation of first N natural numbers and an integer M ≤ N. The task is to find the number of sub-arrays such that the median of the sequence is M.
The median of a sequence is the value of the element which is in the middle of the sequence after sorting it in non-decreasing order. If the length of the sequence is even, the left of two middle elements is used.

Examples:

Input: a[] = { 2, 4, 5, 3, 1}, M = 4
Output: 4
Required sub-arrays are {2, 4, 5}, {4}, {4, 5} and {4, 5, 3}.

Input: a[] = { 1, 2, 3, 4, 5}, M = 5
Output: 1



Approach: The segment p[l..r] has median equals M if and only if M belongs to it and less = greater or less = greater – 1, where less is number of elements in p[l..r] that are strictly less than M and greater is number of elements in p[l..r] that are strictly greater than M. Here we’ve used a fact that p is a permutation (on p[l..r] there is exactly one occurrence of M).

In other words, M belongs to p[l..r] and the value greater – less equals to 0 or 1.

Calculate prefix sums sum[0..n], where sum[i] the value greater-less on the prefix of the length i (i.e., on the subarray p[0..i-1]). For fixed value r it is easy to calculate the number of such l that p[l..r] is suitable. At first, check that M met on [0..r]. Valid values l are such indices that: no M on [0..l-1] and sum[l]=sum[r] or sum[r]=sum[l]+1.

Let’s maintain a number of prefix sums sum[i] to the left of M for each value. We can use just a map c, where c[s] is a number of such indices l that sum[l]=s and l are to the left of m.

So for each r that p[0..r] contains m do ans += c[sum] + c[sum – 1], where sum is the current value greater-less.

Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation of the approach
#include <bits/stdc++.h>
using namespace std;
  
// Function to return the count of sub-arrays
// in the given permutation of first n natural
// numbers such that their median is m
int segments(int n, int p[], int m)
{
    map<int, int> c;
    c[0] = 1;
    bool has = false;
    int sum = 0;
    long long ans = 0;
    for (int r = 0; r < n; r++) {
  
        // If element is less than m
        if (p[r] < m)
            sum--;
  
        // If element greater than m
        else if (p[r] > m)
            sum++;
  
        // If m is found
        if (p[r] == m)
            has = true;
  
        // Count the answer
        if (has)
            ans += c[sum] + c[sum - 1];
  
        // Increment sum
        else
            c[sum]++;
    }
  
    return ans;
}
  
// Driver code
int main()
{
    int a[] = { 2, 4, 5, 3, 1 };
    int n = sizeof(a) / sizeof(a[0]);
    int m = 4;
    cout << segments(n, a, m);
  
    return 0;
}

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# Python3 implementation of the approach
  
# Function to return the count of sub-arrays
# in the given permutation of first n natural
# numbers such that their median is m
def segments(n, p, m):
  
    c = dict()
  
    c[0] = 1
  
    has = False
  
    Sum = 0
  
    ans = 0
  
    for r in range(n):
  
        # If element is less than m
        if (p[r] < m):
            Sum -= 1
  
        # If element greater than m
        elif (p[r] > m):
            Sum += 1
  
        # If m is found
        if (p[r] == m):
            has = True
  
        # Count the answer
        if (has):
            if(Sum in c.keys()):
                ans += c[Sum]
            if Sum-1 in c.keys():
                ans += c[Sum - 1
  
        # Increment Sum
        else:
            c[Sum] = c.get(Sum, 0) + 1
  
    return ans
  
# Driver code
a = [2, 4, 5, 3, 1]
n = len(a)
m = 4
print(segments(n, a, m))
  
# This code is contributed by mohit kumar

chevron_right


PHP

filter_none

edit
close

play_arrow

link
brightness_4
code

<?php
// PHP implementation of the approach 
  
// Function to return the count of sub-arrays 
// in the given permutation of first n natural 
// numbers such that their median is m 
function segments($n, $p, $m
    $c = array(); 
    $c[0] = 1; 
      
    $has = false; 
    $sum = 0; 
    $ans = 0; 
      
    for ($r = 0; $r < $n; $r++) 
    
  
        // If element is less than m 
        if ($p[$r] < $m
            $sum--; 
  
        // If element greater than m 
        else if ($p[$r] > $m
            $sum++; 
  
        // If m is found 
        if ($p[$r] == $m
            $has = true; 
  
        // Count the answer 
        if ($has
            $ans += $c[$sum] + $c[$sum - 1]; 
  
        // Increment sum 
        else
            $c[$sum]++; 
    
  
    return $ans
  
// Driver code 
$a = array( 2, 4, 5, 3, 1 ); 
$n = count($a);
$m = 4; 
  
echo segments($n, $a, $m); 
  
// This code is contributed by Ryuga
?>

chevron_right


Output:

4


My Personal Notes arrow_drop_up

pawanasipugmailcom

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.



Improved By : mohit kumar 29, AnkitRai01