使用另一个数组中的值选择numpy矩阵的行，并计算每列的平均值

Question

假设有一个形状为(r, c)的NumPy数组M，例如：

M = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9],
              [10, 11, 12],
              [13, 14, 15]])  # r = 5; c = 3

有一个长度为r的一维数组a，其中包含0到k-1之间可变的整数，例如：

a = np.array([0, 0, 2, 1, 0])  # k = 4

我想使用a中的值来选择M中的行，以获得如下中间结果：

array([
       [[1, 2, 3], [4, 5, 6], [13, 14, 15]]  # rows of M where a == 0
       [[10, 11, 12]],                       # rows of M where a == 1
       [[7, 8, 9]]                           # rows of M where a == 2
       []                                    # rows of M where a == 3 (there are none)
      ]) 

（我不需要这个中间数组，但只是为了说明而展示它。）返回的结果将是一个 (k, c) 数组，其中包含从此数组中按列计算的平均值：

array([[ 6.,  7.,  8.],   # means where a == 0
   [10., 11., 12.],       # means where a == 1
   [ 7.,  8.,  9.],       # etc.
   [nan, nan, nan]])

我可以使用以下方式完成此操作

np.array([M[a == i].mean(axis=0) for i in range(k)])

但是是否有一种方法（对于大的r和k来说更快），纯粹使用numpy方法而不是使用for循环来创建一个列表（然后必须将其转换回数组）？

- Stuart

---

你不能像你描述的那样拥有一个非矩形数组，但你可以将最终输出作为数组。 - Ehsan

你想使用Pandas还是纯Numpy？ - Ehsan

@Ehsan 谢谢，我已经编辑了问题 - 我不需要非矩形中间数组。 - Stuart

@Ehsan 我正在寻找一种numpy方法来优化速度。 - Stuart

哦，我已经发布了一个关于pandas的。如果没有numpy的话，我会稍后发布一个。我建议为了速度尝试使用循环的numba，可能会更快。 - Ehsan

Answer 1

一个Pandas解决方案（感谢@Quang提供更好的代码建议）：

pd.DataFrame(M).groupby(a).mean().reindex(np.arange(k)).values
#[[ 6.  7.  8.]
# [10. 11. 12.]
# [ 7.  8.  9.]
# [nan nan nan]]

使用 @Divakar 的 benchit 进行比较：

#@Ehsan's solution
def m1(M, a):
  return pd.DataFrame(M).groupby(a).mean().reindex(np.arange(k)).values

#@Divakar's solution 1
def m2(M, a):
  m = a == np.arange(k)[:,None] #or np.equal.outer(np.arange(k), a)
  return m.dot(M)/m.sum(1,keepdims=True)

#@Mad Physicist's solution 1
def m3(M, a):
    index = np.argsort(a)
    a = a[index]
    M = M[index]
    split = np.flatnonzero(np.diff(np.r_[-1, a]))
    means = np.add.reduceat(M, split, axis=0) / np.diff(np.r_[split, a.size])[:, None]
    result = np.full((k, M.shape[1]), np.nan)
    result[a[split], :] = means
    return result

#@Mad Physicist's solution 2
def m4(M, a):
    u, ind, cnt = np.unique(a, return_inverse=True, return_counts=True)
    ind = np.argsort(ind)
    split = np.cumsum(np.r_[0, cnt[:-1]])
    result = np.full((k, M.shape[1]), np.nan)
    result[u, :] = np.add.reduceat(M[ind], split, axis=0) / cnt[:, None]
    return result

#@Divakar's solution 2
def m5(M, a):
  n = M.shape[1]
  out = np.empty((k,n))
  for i in range(n):
      out[:,i] = np.bincount(a,M[:,i], minlength=k) 
  out /= np.bincount(a,minlength=k)[:,None]
  return out

#@OP's baseline
def m6(M, a):
  return np.array([M[a == i].mean(axis=0) for i in range(k)])

funcs = [m1,m2,m3,m4,m5,m6]
in_ = {n:[np.random.randint(100,size=(n,10)), np.random.randint(0,k,size=n)] for n in [100,1000,10000,100000]}

我认为，根据聚类的数量、列和行的数量，可以采用不同的方法使其更有效率。

当k=10时：

当k=100时：

当k=1000时：

当k=100 且 in_ = {n:[np.random.randint(100,size=(n,1000)), np.random.randint(0,k,size=n)] for n in [100,1000,10000]} ( M 具有1000个较小范围的列) 时：

- Ehsan

---

pd.DataFrame(M).groupby(a).mean().reindex(np.arange(k)) - Quang Hoang

你不需要为了之后删除它而创建一个新的分组列。 - Quang Hoang

我又添加了一个。对于优雅的解决方案和始终受欢迎的基准测试，加1分。 - Mad Physicist

非常感谢您的基准测试！您能否添加我的原始解决方案进行比较？m5函数没有返回任何内容 - 这可能会影响基准测试吗？此外，我注意到m5在M中的列上使用for循环，因此在更宽的矩阵上可能表现不佳。 - Stuart

@Stuart 没关系，我认为这不会有太大的影响。 - Ehsan

Answer 2

方法一

使用基于BLAS的矩阵乘法的NumPy解决方案 -

In [46]: m = a == np.arange(k)[:,None] #or np.equal.outer(np.arange(k), a)

In [47]: m.dot(M)/m.sum(1,keepdims=True)
Out[47]: 
array([[ 6.,  7.,  8.],
       [10., 11., 12.],
       [ 7.,  8.,  9.],
       [nan, nan, nan]])

为了提高大型数组的性能，请将掩码转换为浮点数，并使用np.bincount来计算掩码元素的数量 -

In [108]: m.astype(float).dot(M)/np.bincount(a,minlength=k)[:,None]
Out[108]: 
array([[ 6.,  7.,  8.],
       [10., 11., 12.],
       [ 7.,  8.,  9.],
       [nan, nan, nan]])

方法二

另一种方法是使用np.bincount在列之间应用 -

n = M.shape[1]
out = np.empty((k,n))
for i in range(n):
    out[:,i] = np.bincount(a,M[:,i], minlength=k) 
out /= np.bincount(a,minlength=k)[:,None]

基准测试

对@Ehsan帖子中列出的所有方法进行基准测试。

使用benchit包（打包了一些基准测试工具；免责声明：我是其作者）来对提议的解决方案进行基准测试：

import benchit
funcs = [m1,m2,m3,m4,m5,m6]
in_ = {(nr,nc,k):[np.random.randint(100,size=(nr,nc)), np.random.randint(0,k,size=nr), k] for nr in [100,1000,10000,100000] for nc in [10,100,200] for k in [10, 100, 1000]}
t = benchit.timings([m1, m2, m3, m4, m5, m6], in_, multivar=True, input_name=['nrows','ncols','k'])
t.plot(logx=True, save='timings.png')

因为 m6 是原始的。让我们相对于它加速：

t.speedups(m6).plot(logx=True, logy=True, save='speedups.png')

在所有数据集中，没有明确的优胜者。然而，m5 在小的 ncols 上表现良好，而 m1 在较大的 ncols 上表现更好。

- Divakar

---

接受因为对于许多矩阵长度和k值来说，方法2似乎是最快的选择，只要矩阵不是非常宽。 - Stuart

@Stuart 添加了一些详细的绘图以使事情更清晰。在所有变化中仍然没有明显的赢家。 - Divakar

Answer 3

一个纯粹的numpy解决方案需要找到将M的行放入组中的排序顺序。例如:

index = np.argsort(a)

接下来您可以找到分割点：

split = np.flatnonzero(np.diff(np.r_[-1, a[index]]))

然后使用np.add.reduceat将结果相加，以获得总和：

sums = np.add.reduceat(M[index], split, axis=0)

另外一个 diff 给出了因子：

lengths = np.diff(np.r_[split, a.size])

你可以使用a中编码的索引填充结果：

result = np.full((k, M.shape[1]), np.nan)
result[a[index][split], :] = sums / lengths[:, None]

总之：

def m3(M, a, k):
    index = np.argsort(a)
    a = a[index]
    M = M[index]
    split = np.flatnonzero(np.diff(np.r_[-1, a]))
    means = np.add.reduceat(M, split, axis=0) / np.diff(np.r_[split, a.size])[:, None]
    result = np.full((k, M.shape[1]), np.nan)
    result[a[split], :] = means
    return result

解决方案可以使用 np.unique 更精简，它会为您预先计算拆分点和索引。 您需要对索引进行 argsort 反转，然后应用累积总和以获得正确的拆分点：

def m4(M, a, k):
    u, ind, cnt = np.unique(a, return_inverse=True, return_counts=True)
    ind = np.argsort(ind)
    split = np.cumsum(np.r_[0, cnt[:-1]])
    result = np.full((k, M.shape[1]), np.nan)
    result[u, :] = np.add.reduceat(M[ind], split, axis=0) / cnt[:, None]
    return result

- Mad Physicist