dado un conjunto de enteros como
[1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5]
Necesito enmascarar elementos que se repiten más de N
veces. Para aclarar: el objetivo principal es recuperar la matriz de máscara booleana, para usarla más adelante para los cálculos de agrupamiento.
Se me ocurrió una solución bastante complicada
import numpy as np
bins = np.array([1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5])
N = 3
splits = np.split(bins, np.where(np.diff(bins) != 0)[0]+1)
mask = []
for s in splits:
if s.shape[0] <= N:
mask.append(np.ones(s.shape[0]).astype(np.bool_))
else:
mask.append(np.append(np.ones(N), np.zeros(s.shape[0]-N)).astype(np.bool_))
mask = np.concatenate(mask)
dando por ejemplo
bins[mask]
Out[90]: array([1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5])
¿Hay una mejor manera de hacer esto?
EDITAR, # 2
¡Muchas gracias por las respuestas! Aquí hay una versión delgada de la trama de referencia de MSeifert. Gracias por señalarme simple_benchmark
. Mostrando solo las 4 opciones más rápidas:
Conclusión
La idea propuesta por Florian H , modificada por Paul Panzer parece ser una excelente manera de resolver este problema, ya que es bastante simple y directo numpy
. numba
Sin embargo, si está de acuerdo con el uso , la solución de MSeifert supera a la otra.
Elegí aceptar la respuesta de MSeifert como solución, ya que es la respuesta más general: maneja correctamente matrices arbitrarias con bloques (no únicos) de elementos repetidos consecutivos. En caso de numba
que no se pueda , ¡ la respuesta de Divakar también vale la pena!