Open In App

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:
The required sub-arrays are {2, 4, 5}, {4}, {4, 5} and {4, 5, 3}.
 

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

Approach: The segment p[l..r] has a median equals M if and only if M belongs to it and less = greater or less = greater – 1, where less is the number of elements in p[l..r] that are strictly less than M and greater is a 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 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 the fixed value r it is easy to calculate the number of so l that p[l..r] is suitable. 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 keep a number of prefix sums sum[i] to the left of M for each value. We can just use a map c, where c[s] is a number of 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++ 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;
}




// Java implementation of the approach
import java.util.HashMap;
 
class GFG
{
 
    // Function to return the count of sub-arrays
    // in the given permutation of first n natural
    // numbers such that their median is m
    public static int segments(int n, int[] p, int m)
    {
        HashMap<Integer, Integer> c = new HashMap<>();
        c.put(0, 1);
        boolean has = false;
        int sum = 0;
        int 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.get(sum) == null ? 0 :
                        c.get(sum)) +
                       (c.get(sum - 1) == null ? 0 :
                        c.get(sum - 1));
 
            // Increment sum
            else
                c.put(sum, c.get(sum) == null ? 1 :
                           c.get(sum) + 1);
        }
        return ans;
    }
 
    // Driver code
    public static void main(String[] args)
    {
        int[] a = { 2, 4, 5, 3, 1 };
        int n = a.length;
        int m = 4;
        System.out.println(segments(n, a, m));
    }
}
 
// This code is contributed by
// sanjeev2552




# 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




// C# implementation of the approach
using System;
using System.Collections.Generic;            
 
class GFG
{
 
    // Function to return the count of sub-arrays
    // in the given permutation of first n natural
    // numbers such that their median is m
    public static int segments(int n, int[] p, int m)
    {
        Dictionary<int, int> c = new Dictionary<int, int>();
        c.Add(0, 1);
        bool has = false;
        int sum = 0;
        int 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.ContainsKey(sum) ? 0 :
                         c[sum]) +
                    (!c.ContainsKey(sum - 1) ? 0 :
                      c[sum - 1]);
 
            // Increment sum
            else
                c.Add(sum, !c.ContainsKey(sum) ? 1 :
                            c[sum] + 1);
        }
        return ans;
    }
 
    // Driver code
    public static void Main(String[] args)
    {
        int[] a = { 2, 4, 5, 3, 1 };
        int n = a.Length;
        int m = 4;
        Console.WriteLine(segments(n, a, m));
    }
}
 
// This code is contributed by 29AjayKumar




<?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
?>




<script>
 
// javascript 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)
{
    var c = new Map();
    c.set(0,1);
    var hs = false;
    var sum = 0;
    var ans = 0;
    var r;
    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)
            hs = true;
 
        // Count the answer
        if (hs){
            if(c.has(sum) && c.has(sum-1))
              ans += c.get(sum) + c.get(sum - 1);
            else if(c.has(sum))
              ans += c.get(sum);
            else if(c.has(sum-1))
             ans += c.get(sum-1);
        }
 
        // Increment sum
        else{
            if(c.has(sum))
             c.set(sum,c.get(sum)+1);
            else
              c.set(sum,1);
        }
    }
 
    return ans;
}
 
// Driver code
 
    var a = [2, 4, 5, 3, 1];
    var n = a.length;
    var m = 4;
    document.write(segments(n, a, m));
 
</script>

Output: 
4

 

Time Complexity: O(N)
Auxiliary Space: O(N)


Article Tags :