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

Question

给定一个 numpy 数组 M形状(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)])

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

Answer 1

纯 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