dada uma matriz de números inteiros como
[1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5]
Eu preciso mascarar elementos que se repetem mais que N
vezes. Para esclarecer: o objetivo principal é recuperar a matriz de máscaras booleanas e usá-la posteriormente para cálculos de classificação.
Eu vim com uma solução 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 exemplo
bins[mask]
Out[90]: array([1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5])
Existe uma maneira melhor de fazer isso?
EDIT, # 2
Muito obrigado pelas respostas! Aqui está uma versão reduzida do gráfico de benchmark do MSeifert. Obrigado por me indicar simple_benchmark
. Mostrando apenas as 4 opções mais rápidas:
Conclusão
A idéia proposta por Florian H , modificada por Paul Panzer, parece ser uma ótima maneira de resolver esse problema, pois é bem simples e direto numpy
. numba
No entanto, se você estiver bem em usar , a solução do MSeifert supera a outra.
Optei por aceitar a resposta do MSeifert como solução, pois é a resposta mais geral: lida corretamente com matrizes arbitrárias com blocos (não exclusivos) de elementos repetidos consecutivos. Caso numba
não seja possível, a resposta de Divakar também vale a pena dar uma olhada!