删除具有比给定数量更多特定值的行的最佳方法（numpy）

Question

我创建了一个列大小为d 的数字数组，其中的数字是set(p) 的不同组合。接下来，我想删除比p 中的数字实例更多的行。我可以用for 循环来做到这一点，这是我的代码：

import numpy as np
import itertools as it

from collections import Counter

p = [0,0,0,1,1,1,1,2,2,2,2,2,3,3,3,4,4,4,4,4,5,5,5]
assert (len(set(p)) != 1)
cnt = Counter(p)
pos = np.array(list(it.product(set(p), repeat=d)))
ds = []
for i in range(len(pos)):
    for j, k in cnt.items():
        if len(np.where(pos[i] == j)[0]) > k:
            ds.append(i)
pos = np.delete(pos, ds, axis=0)

我正在寻找一种更快的方法来做到这一点。谢谢！

Answer 1

对于正整数，我们可以使用bincount。因此，我们将使用2D bincount solution 获取pos 中每行的分箱计数，并与p 的分箱计数进行比较，并使用该掩码从pos 中选择有效行，就像这样 -

# https://stackoverflow.com/a/46256361/ @Divakar
def bincount2D_vectorized(a):    
    N = a.max()+1
    a_offs = a + np.arange(a.shape[0])[:,None]*N
    return np.bincount(a_offs.ravel(), minlength=a.shape[0]*N).reshape(-1,N)

pos_out = pos[(bincount2D_vectorized(pos)>np.bincount(p)).any(1)]