Recuento de subarreglos de un arreglo dado con una mediana de al menos X

Dada una array arr[] de enteros con longitud N y un entero X , la tarea es calcular el número de subarreglos con una mediana mayor o igual que el entero X dado .

Ejemplos:

Entrada: N=4, A = [5, 2, 4, 1], X = 4
Salida: 7
Explicación: Para el subarreglo [5], la mediana es 5. (>= 4)
Para el subarreglo [5, 2], la mediana es 5. (>= 4)
Para el subarreglo [5, 2, 4], la mediana es 4. (>= 4)
Para el subarreglo [5, 2, 4, 1], la mediana es 4. (>= 4)
Para el subarreglo [2, 4], la mediana es 4. (>= 4)
Para el subarreglo [4], la mediana es 4. (>= 4)
Para el subarreglo [4, 1], la mediana es 4. (>= 4)

Entrada: N = [3, 7, 2, 0, 1, 5], X = 10
Salida: 0
Explicación: No hay subarreglos con una mediana mayor o igual a X.

Planteamiento:  El problema se puede resolver con base en la siguiente idea.

Para encontrar un subarreglo con una mediana mayor o igual a X, al menos la mitad de los elementos deben ser mayores o iguales a X.

Siga los pasos a continuación para implementar la idea anterior:

• Reemplace cada elemento de una array con 1 si es mayor o igual a X , de lo contrario, reemplácelo con -1 .
• Con base en la idea anterior, para que el nuevo arreglo, la mediana de cualquier subarreglo sea mayor o igual a X , su suma de elementos debe ser mayor o igual a 0.
• Para calcular el número de subarreglo con una suma mayor o igual a 0:
  • Encuentre la suma de prefijos para cada índice de la nueva array.
  • Recorra la array de prefijos recién creada a partir del índice 1 y calcule la cantidad de elementos anteriores con un valor menor o igual que el valor actual.
  • Agregue todos aquellos en la respuesta final, ya que también formarán un subarreglo con el actual que satisface todas las condiciones.
  • Después de encontrarlo para un índice, agregue el valor actual a un conjunto múltiple .
• Devuelve la respuesta final.

Nota : Para calcular de manera eficiente la cantidad de elementos con un valor menor o igual a Y , use estructuras de datos basadas en políticas .

A continuación se muestra la implementación del enfoque anterior:

// C++ code to implement the above approach
 
#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;
 
// A new data structure defined.
typedef tree<int, null_type, less_equal<int>, rb_tree_tag,
             tree_order_statistics_node_update>
    ordered_set;
 
// Function to to find the Number of
// subarrays  with median greater than
// or equal to X.
long long findNumberOfSubarray(int arr[],
                               int n, int X)
{
    // Build new array by comparing it with 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;
        }
    }
 
    // Build new array in which
    // at i-th index, Sum of first i elements
    // are stored
    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];
    }
 
    // Store the answer
    // Using long long because
    // it can exceed the storage limit of int
    long long ans = 0;
 
    // For storing already traversed values
    ordered_set s;
    s.insert(0);
 
    // Iterating forwards from 0 to n-1.
    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;
}
 
// Driver Code
int main()
{
    int N = 4, X = 4;
    int arr[] = { 5, 2, 4, 1 };
 
    // Function call
    long long ans
        = findNumberOfSubarray(arr, N, X);
    cout << ans;
    return 0;
}

# python3 code to implement the above approach
import bisect
 
# Function to to find the Number of
# subarrays with median greater than
# or equal to X.
def findNumberOfSubarray(arr, n, X):
 
    # Build new array by comparing it with 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
 
    # Build new array in which
    # at i-th index, Sum of first i elements
    # are stored
    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]
 
    # Store the answer
    # Using long long because
    # it can exceed the storage limit of int
    ans = 0
 
    # For storing already traversed values
    s = set()
    s.add(0)
 
    # Iterating forwards from 0 to n-1.
    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
 
# Driver Code
if __name__ == "__main__":
 
    N, X = 4, 4
    arr = [5, 2, 4, 1]
 
    # Function call
    ans = findNumberOfSubarray(arr, N, X)
    print(ans)
 
    # This code is contributed by rakeshsahni

7

Complejidad temporal: O(N * logN)
Espacio auxiliar: O(N)