python - numpy 一维数组 : mask elements that repeat more than n times

Question

问:给定一个整数数组，例如

[1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5]

我需要屏蔽重复次数超过 N 的元素次。目标是检索 bool 掩码数组。
我想出了一个相当复杂的解决方案:

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])

有没有更好的方法来做到这一点？

总结:这是 MSeifert 基准图的精简版(感谢您指向 simple_benchmark)。显示四个性能最高的选项:

Florian H 提出的想法, 由 Paul Panzer 修改似乎是解决这个问题的好方法，因为它非常简单， numpy -只要。如果您可以使用 numba , MSeifert's solution胜过对方。
我选择接受 MSeifert 的答案作为解决方案，因为它是更一般的答案:它正确处理具有(非唯一)连续重复元素 block 的任意数组。万一 numba是不行的， Divakar's answer也值得一看。

Answer 1

免责声明:这只是@FlorianH 想法的更合理的实现:

def f(a,N):
    mask = np.empty(a.size,bool)
    mask[:N] = True
    np.not_equal(a[N:],a[:-N],out=mask[N:])
    return mask

对于更大的数组，这有很大的不同:

a = np.arange(1000).repeat(np.random.randint(0,10,1000))
N = 3

print(timeit(lambda:f(a,N),number=1000)*1000,"us")
# 5.443050000394578 us

# compare to
print(timeit(lambda:[True for _ in range(N)] + list(bins[:-N] != bins[N:]),number=1000)*1000,"us")
# 76.18969900067896 us