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
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: 1
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> |
4
Time Complexity: O(N)
Auxiliary Space: O(N)