Python 添加具有相同bin分配的numpy数组元素/切片

我有一些数组

A

，数组

bin

的对应元素包含每行的bin分配。我想构造一个数组

S

，这样

S[0, :] = (A[(bins == 0), :]).sum(axis=0)

这对于

np.stack

和列表理解来说相当容易，但它似乎过于复杂，可读性也不太好。是否有一种更通用的方法来求和（甚至应用一些通用函数）具有bin赋值的数组切片

scipy.stats.binned_statistic

是正确的，但要求计算函数的bin赋值和值是相同的形状（因为我使用的是切片，所以情况并非如此）

例如，如果

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

及

那么它应该会导致

S = np.array([[10., 10., 10., 10.],
              [2.,  3.,  4.,  5. ],
              [8.,  7.,  6.,  5. ]])

---

您可以使用

np.add.reduceat

：

import numpy as np
# index to sort the bins
sort_index = bins.argsort()

# indices where the array needs to be split at
indices = np.concatenate(([0], np.where(np.diff(bins[sort_index]))[0] + 1))

# sum values where the bins are the same
np.add.reduceat(A[sort_index], indices, axis=0)

# array([[ 10.,  10.,  10.,  10.],
#        [  2.,   3.,   4.,   5.],
#        [  8.,   7.,   6.,   5.]])

---

下面是一种使用

矩阵乘法的方法-

样本运行-

In [40]: A = np.array([[1., 2., 3., 4.],
    ...:               [2., 3., 4., 5.],
    ...:               [9., 8., 7., 6.],
    ...:               [8., 7., 6., 5.]])

In [41]: bins = np.array([0, 1, 0, 2])

In [42]: (bins == np.arange(bins.max()+1)[:,None]).dot(A)
Out[42]: 
array([[ 10.,  10.,  10.,  10.],
       [  2.,   3.,   4.,   5.],
       [  8.,   7.,   6.,   5.]])

性能提升

创建掩码
（bins==np.arange（bins.max（）+1）[：，None]）
的更有效方法如下-

mask = np.zeros((bins.max()+1, len(bins)), dtype=bool)
mask[bins, np.arange(len(bins))] = 1

这比@Psidom的解决方案快30%，所以接受这个。这对我来说有点直截了当，但两者都有效。
In [40]: A = np.array([[1., 2., 3., 4.],
    ...:               [2., 3., 4., 5.],
    ...:               [9., 8., 7., 6.],
    ...:               [8., 7., 6., 5.]])

In [41]: bins = np.array([0, 1, 0, 2])

In [42]: (bins == np.arange(bins.max()+1)[:,None]).dot(A)
Out[42]: 
array([[ 10.,  10.,  10.,  10.],
       [  2.,   3.,   4.,   5.],
       [  8.,   7.,   6.,   5.]])

mask = np.zeros((bins.max()+1, len(bins)), dtype=bool)
mask[bins, np.arange(len(bins))] = 1