Count of Subarrays of given Array with median at least X
Last Updated :
09 Mar, 2023
Given an array arr[]of integers with length N and an integer X, the task is to calculate the number of subarrays with median greater than or equal to the given integer X.
Examples:
Input: N=4, A = [5, 2, 4, 1], X = 4
Output: 7
Explanation: For subarray [5], median is 5. (>= 4)
For subarray [5, 2], median is 5. (>= 4)
For subarray [5, 2, 4], median is 4. (>= 4)
For subarray [5, 2, 4, 1], median is 4. (>= 4)
For subarray [2, 4], median is 4. (>= 4)
For subarray [4], median is 4. (>= 4)
For subarray [4, 1], median is 4. (>= 4)
Input: N = [3, 7, 2, 0, 1, 5], X = 10
Output: 0
Explanation: There are no subarrays with median greater than or equal to X.
Approach: The problem can be solved based on the following idea.
To find a subarray with median greater or equal to X at least half of the elements should be greater than or equal to X.
Follow the below steps to implement the above idea:
- Replace each element of an array with 1 if it is greater than or equal to X, else replace it with -1.
- Based on the above idea, for the new array, median of any subarray to be greater than or equal to X, its sum of elements should be greater than or equal to 0.
- For calculating the number of subarray with a sum greater than or equal to 0:
- Find prefix sum up to each index of the new array.
- Traverse the newly created prefix array starting from index 1 and calculate the number of elements before it with a value less than or equal to the current value.
- Add all those in the final answer as they will also form a subarray with the current one satisfying all conditions.
- After finding it for an index, add the current value to a multiset.
- Return the final answer.
Note: For efficiently calculating the number of elements with a value less than or equal to Y, use policy-based data structures.
Below is the implementation of the above approach:
C++
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <functional>
#include <iostream>
using namespace __gnu_pbds;
using namespace std;
typedef tree< int , null_type, less_equal< int >, rb_tree_tag,
tree_order_statistics_node_update>
ordered_set;
long long findNumberOfSubarray( int arr[],
int n, int X)
{
int new_array[n];
for ( int i = 0; i < n; i++) {
if (arr[i] >= X) {
new_array[i] = 1;
}
else {
new_array[i] = -1;
}
}
int pref_sum[n];
pref_sum[0] = new_array[0];
for ( int i = 1; i < n; i++) {
pref_sum[i] = pref_sum[i - 1]
+ new_array[i];
}
long long ans = 0;
ordered_set s;
s.insert(0);
for ( int i = 0; i < n; i++) {
int less_than
= s.order_of_key(pref_sum[i] + 1);
ans += less_than;
s.insert(pref_sum[i]);
}
return ans;
}
int main()
{
int N = 4, X = 4;
int arr[] = { 5, 2, 4, 1 };
long long ans
= findNumberOfSubarray(arr, N, X);
cout << ans;
return 0;
}
|
Java
import java.util.*;
class GFG {
public static int findNumberOfSubarray( int [] arr, int n, int X) {
int [] newArray = new int [n];
for ( int i = 0 ; i < n; i++) {
if (arr[i] >= X) {
newArray[i] = 1 ;
} else {
newArray[i] = - 1 ;
}
}
int [] prefSum = new int [n];
prefSum[ 0 ] = newArray[ 0 ];
for ( int i = 1 ; i < n; i++) {
prefSum[i] = prefSum[i - 1 ] + newArray[i];
}
int ans = 0 ;
HashSet<Integer> set = new HashSet<>();
set.add( 0 );
for ( int i = 0 ; i < n; i++) {
int lessThan = Collections.binarySearch( new ArrayList<>(set), prefSum[i] + 1 );
if (lessThan < 0 ) {
lessThan = -(lessThan + 1 );
}
if (set.contains(prefSum[i] + 1 )) {
lessThan++;
}
ans += lessThan;
set.add(prefSum[i]);
}
return ans;
}
public static void main(String[] args) {
int N = 4 , X = 4 ;
int [] arr = { 5 , 2 , 4 , 1 };
int ans = findNumberOfSubarray(arr, N, X);
System.out.println(ans);
}
}
|
Python3
import bisect
def findNumberOfSubarray(arr, n, X):
new_array = [ 0 for _ in range (n)]
for i in range ( 0 , n):
if (arr[i] > = X):
new_array[i] = 1
else :
new_array[i] = - 1
pref_sum = [ 0 for _ in range (n)]
pref_sum[ 0 ] = new_array[ 0 ]
for i in range ( 1 , n):
pref_sum[i] = pref_sum[i - 1 ] + new_array[i]
ans = 0
s = set ()
s.add( 0 )
for i in range ( 0 , n):
less_than = bisect.bisect_left(
sorted (s), pref_sum[i] + 1 , lo = 0 , hi = len (s))
if pref_sum[i] + 1 in s:
less_than + = 1
ans + = less_than
s.add(pref_sum[i])
return ans
if __name__ = = "__main__" :
N, X = 4 , 4
arr = [ 5 , 2 , 4 , 1 ]
ans = findNumberOfSubarray(arr, N, X)
print (ans)
|
C#
using System;
using System.Collections.Generic;
class GFG
{
public static int findNumberOfSubarray( int [] arr, int n, int X)
{
int [] newArray = new int [n];
for ( int i = 0; i < n; i++)
{
if (arr[i] >= X)
{
newArray[i] = 1;
}
else
{
newArray[i] = -1;
}
}
int [] prefSum = new int [n];
prefSum[0] = newArray[0];
for ( int i = 1; i < n; i++)
{
prefSum[i] = prefSum[i - 1] + newArray[i];
}
int ans = 0;
HashSet< int > set = new HashSet< int >();
set .Add(0);
for ( int i = 0; i < n; i++)
{
int lessThan = Array.BinarySearch( new List< int >( set ).ToArray(), prefSum[i] + 1);
if (lessThan < 0)
{
lessThan = -(lessThan + 1);
}
if ( set .Contains(prefSum[i] + 1))
{
lessThan++;
}
ans += lessThan;
set .Add(prefSum[i]);
}
return ans;
}
static void Main( string [] args)
{
int N = 4, X = 4;
int [] arr = { 5, 2, 4, 1 };
int ans = findNumberOfSubarray(arr, N, X);
Console.WriteLine(ans);
}
}
|
Javascript
function findNumberOfSubarray(arr, n, X) {
let new_array = Array(n).fill(0);
for (let i = 0; i < n; i++) {
if (arr[i] >= X) {
new_array[i] = 1;
} else {
new_array[i] = -1;
}
}
let pref_sum = Array(n).fill(0);
pref_sum[0] = new_array[0];
for (let i = 1; i < n; i++) {
pref_sum[i] = pref_sum[i - 1] + new_array[i];
}
let ans = BigInt(0);
let s = new Set();
s.add(0);
for (let i = 0; i < n; i++) {
let sortedS = Array.from(s).sort( function (a, b) {
return a - b
});
let less_than = bisect_left(sortedS, pref_sum[i] + 1, 0, sortedS.length);
if (s.has(pref_sum[i] + 1)) {
less_than += 1;
}
ans += BigInt(less_than);
s.add(pref_sum[i]);
}
return ans;
}
function bisect_left(sortedArr, x, lo, hi) {
lo = lo || 0;
hi = hi || sortedArr.length;
while (lo < hi) {
let mid = (lo + hi) >>> 1;
if (sortedArr[mid] < x) {
lo = mid + 1;
} else {
hi = mid;
}
}
return lo;
}
let N = 4,
X = 4;
let arr = [5, 2, 4, 1];
let ans = findNumberOfSubarray(arr, N, X);
console.log(ans);
|
Time Complexity: O(N * logN)
Auxiliary Space: O(N)
Share your thoughts in the comments
Please Login to comment...