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++
// 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 mint 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 codeint 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
// 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
# 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 mdef 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 codea = [2, 4, 5, 3, 1]n = len(a)m = 4print(segments(n, a, m))# This code is contributed by mohit kumar |
C#
// 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// 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?> |
Javascript
<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 mfunction 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)
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
