在numpy数组中每行提取固定数量的元素

Question

假设我有一个数组a，和一个布尔数组b，我想从a的每一行的有效元素中提取固定数量的元素。有效元素是由 b.

指示的元素

这是一个例子：

a = np.arange(24).reshape(4,6)
b = np.array([[0,0,1,1,0,0],[0,1,0,1,0,1],[0,1,1,1,1,0],[0,0,0,0,1,1]]).astype(bool)
x = []
for i in range(a.shape[0]):
    c = a[i,b[i]]
    d = np.random.choice(c, 2)
    x.append(d)

这里我使用了一个for循环，如果这些数组又大又高维的话会很慢。有没有更有效的方法来做到这一点？谢谢。

Answer 1

生成形状为 a.
将此矩阵乘以掩码 b 将无效元素设置为零。
Select 每行的 k 最大索引（模拟无偏随机 k-仅来自该行中有效元素的样本）。
（可选）使用这些索引获取元素。

a = np.arange(24).reshape(4,6)
b = np.array([[0,0,1,1,0,0],[0,1,0,1,0,1],[0,1,1,1,1,0],[0,0,0,0,1,1]])
k = 2

r = np.random.uniform(size=a.shape)
indices = np.argpartition(-r * b, k)[:,:k]

从索引中获取元素：

>>> indices
array([[3, 2],
       [5, 1],
       [3, 2],
       [4, 5]])
>>> a[np.arange(a.shape[0])[:,None], indices]
array([[ 3,  2],
       [11,  7],
       [15, 14],
       [22, 23]])

在numpy数组中每行提取固定数量的元素

Extract fixed number of elements per row in numpy array

arrays

numpy

numpy-ufunc

numpy-ndarray

array-broadcasting