python - 基于多标签掩码对 numpy 数组求和

Question

假设我有一个数组:

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

还有一个多标签掩码:

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

我的目标是将 x 的条目加在一起，按 labels 分组。例如:

n_labels = np.max(labels) + 1
out = np.empty(n_labels)
for label in range(n_labels):
    mask = labels == label
    out[label] = np.sum(x[mask])

>>> out
np.array([1, 20, 15])

但是，随着 x 和 n_labels 变大，我认为这是低效的。每次迭代，您只是将 x 的条目数的一小部分加在一起，但仍然重新计算所有 labels 的掩码(在表达式 labels == label) 并随后对所有 x 进行索引(在表达式 x[mask] 中)。当 x 和 n_labels 变大时，是否有更有效的方法来做到这一点？

Answer 1

您可以使用 bincount有权重:

np.bincount(labels.ravel(), weights=x.ravel())

输出:

array([ 1., 20., 15.])