Python 具有条件的每行唯一值的二维矢量化

考虑如下所示的数组和函数定义：

import numpy as np

a = np.array([[2, 2, 5, 6, 2, 5],
              [1, 5, 8, 9, 9, 1],
              [0, 4, 2, 3, 7, 9],
              [1, 4, 1, 1, 5, 1],
              [6, 5, 4, 3, 2, 1],
              [3, 6, 3, 6, 3, 6],
              [0, 2, 7, 6, 3, 4],
              [3, 3, 7, 7, 3, 3]])

def grpCountSize(arr, grpCount, grpSize):    
    count = [np.unique(row, return_counts=True)  for row in arr]
    valid = [np.any(np.count_nonzero(row[1] == grpSize) == grpCount) for row in count]
    return valid

该函数的作用是返回数组

a

中的行，这些行具有完全相同的

grpCount

元素组，每个元素组包含完全相同的

grpSize

例如：

# which rows have exactly 1 group that holds exactly 2 identical elements?
out = a[grpCountSize(a, 1, 2)]

正如预期的那样，代码输出

out=[[2,2,5,6,2,5]，[3,3,7,7,3,3]]

。 第一个输出行正好有一组2（即：5,5），而第二个输出行也正好有一组2（即：7,7）

同样地：

# which rows have exactly 2 groups that each hold exactly 3 identical elements?
out = a[grpCountSize(a, 2, 3)]

这将产生

out=[[3,6,3,6,3,6]]

，因为只有此行有两个组，每个组正好包含3个元素（即：3,3,3和6,6,6）

---

问题：我的实际数组只有6列，但它们可能有数百万行。代码完全按照预期工作，但对于长数组来说速度非常慢。有什么方法可以加快速度吗？

np.unique

对数组进行排序，从而降低您的使用效率。使用

np.bincount

，这样很可能会节省一些时间（取决于数组的形状和值）。您也不再需要

np.任何
：

def grpCountSize(arr, grpCount, grpSize):    
    count = [np.bincount(row) for row in arr]
    valid = [np.count_nonzero(row == grpSize) == grpCount for row in count]
    return valid

另一种可能节省更多时间的方法是为所有行使用相同数量的存储箱，并创建一个数组：

def grpCountSize(arr, grpCount, grpSize):
    m = arr.max()
    count = np.stack([np.bincount(row, minlength=m+1) for row in arr])
    return (count == grpSize).sum(1)==grpCount

另一个升级是使用。例如（请注意，上面帖子中测试的Numba解决方案速度更快。我只是提供了numpy解决方案作为示例。您可以使用上面链接的帖子中建议的任何函数替换该函数）：

以上所有解决方案的输出：

a[grpCountSize2(a, 1, 2)]
#array([[2, 2, 5, 6, 2, 5],
#       [3, 3, 7, 7, 3, 3]])

numpy
中的巨大加速来自于对整个数组应用编译方法。你的代码必须一行一行地工作；如果您无法找到一种同时处理多行的方法，那么您必须尝试加快逐行操作的速度。我怀疑您可以将
unique
和
count
操作一起应用于一行，但没有意识到与当前的两个循环相比，than是否会提高时间。另一个选项是使用
numba
，但我不会对此进行探讨。非常清楚和详细的解释！我是Python新手，还没有尝试过实现NUBA，但我认为这对我来说是值得研究的。如果您只是想使用NUBA而不想深入了解其细节，只需在提供的链接中用任何NUBA函数替换
bincount2D\u矢量化的
。
a[grpCountSize2(a, 1, 2)]
#array([[2, 2, 5, 6, 2, 5],
#       [3, 3, 7, 7, 3, 3]])