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

Question

给定一个M形状的 numpy 数组，(r, c)例如

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

和一个a长度为 的一维数组r，包含可以从 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)])

但是有没有一种方法（希望对于 larger和更快k）纯粹使用 numpy 方法而不是使用 for 循环来制作列表（然后必须将其转换回数组）？

Answer 1

熊猫解决方案（感谢@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 列和较小范围的行）

Answer 2

方法#1

利用 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]])

方法#2

另一个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对于小 ncolm1似乎很好，而在较大的 ncol 上似乎更好。

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