учитывая массив целых чисел, таких как
[1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5]
Мне нужно замаскировать элементы, которые повторяются больше, чем N
раз. Чтобы уточнить: основная цель состоит в том, чтобы получить массив логических масок, чтобы использовать его позже для вычисления биннинга.
Я придумал довольно сложное решение
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)
например,
bins[mask]
Out[90]: array([1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5])
Есть ли лучший способ сделать это?
РЕДАКТИРОВАТЬ, № 2
Большое спасибо за ответы! Вот тонкая версия эталонного сюжета MSeifert. Спасибо за указание на меня simple_benchmark
. Показаны только 4 самых быстрых варианта:
Вывод
Идея, предложенная Флорианом Н и модифицированная Полом Панцером, кажется отличным способом решения этой проблемы, поскольку она довольно прямолинейна и numpy
единственна. Однако, если вы в порядке с использованием numba
, решение MSeifert превосходит другие.
Я решил принять ответ MSeifert в качестве решения, так как это более общий ответ: он правильно обрабатывает произвольные массивы с (неуникальными) блоками последовательных повторяющихся элементов. В случае numba
, если нет, ответ Divakar также стоит посмотреть!