Numpy 根据条件拆分数组，无需 for 循环

Question

所以假设我有一个 numpy 数组，它在 2d space 中保存点，如下所示

np.array([[3, 2], [4, 4], [5, 4], [4, 2], [4, 6], [9, 5]])

我还有一个 numpy 数组，将每个点标记为一个数字，这个数组是一个一维数组，长度为点数组中的点数。

np.array([0, 1, 1, 0, 2, 1])

现在我想从标签数组中获取具有索引的每个点的平均值。因此，对于所有标签为 0 的点，取这些点的平均值。我目前解决这个问题的方法如下

return np.array([points[labels==k].mean(axis=0) for k in range(k)])

其中 k 是标签数组中的最大数字，或者称为标记点的方法数。

我想要一种不使用 for 循环的方法，也许我还没有发现一些 numpy 功能？

Answer 1

方法 #1： 我们可以在 braodcasting -

的帮助下利用 matrix-multiplication

mask = labels == np.arange(labels.max()+1)[:,None]
out = mask.dot(points)/np.bincount(labels).astype(float)[:,None]

样本运行-

In [36]: points = np.array([[3, 2], [4, 4], [5, 4], [4, 2], [4, 6], [9, 5]]) 
    ...: labels = np.array([0, 1, 1, 0, 2, 1])

# Original soln
In [37]: L = labels.max()+1

In [38]: np.array([points[labels==k].mean(axis=0) for k in range(L)])
Out[38]: 
array([[3.5       , 2.        ],
       [6.        , 4.33333333],
       [4.        , 6.        ]])

# Proposed soln
In [39]: mask = labels == np.arange(labels.max()+1)[:,None]
    ...: out = mask.dot(points)/np.bincount(labels).astype(float)[:,None]

In [40]: out
Out[40]: 
array([[3.5       , 2.        ],
       [6.        , 4.33333333],
       [4.        , 6.        ]])

方法 #2： 与 np.add.at -

sums = np.zeros((labels.max()+1,points.shape[1]),dtype=float)
np.add.at(sums,labels,points)
out = sums/np.bincount(labels).astype(float)[:,None]

方法#3：如果从0到max-label的序列中的所有数字都出现在labels中，我们也可以使用np.add.reduceat -

sidx = labels.argsort()
sorted_points = points[sidx]
sums = np.add.reduceat(sorted_points,np.r_[0,np.bincount(labels)[:-1].cumsum()])
out = sums/np.bincount(labels).astype(float)[:,None]

Numpy 根据条件拆分数组，无需 for 循环

Numpy split array based on condition without for loop

python

arrays

performance

numpy

vectorization