在 Python 中对矩阵的选定元素求和

Question

我有一个 [nxn] 矩阵，其中包含属于不同组的值，以及一个 [1 xn] 向量定义每个元素属于哪个组。(n 通常 ~1E4, 在这个例子中 n=4)我想计算通过将属于同一组的所有元素相加而获得的矩阵。我使用 np.where() 来计算每个组的元素所在的索引。当我使用计算的索引时，我没有获得预期的元素，因为我选择了位置对而不是范围（我习惯于 Matlab，在那里我可以简单地选择 M(idx1,idx2) ）。import numpy as npn=4M = np.random.rand(n,n)print(M)# This vector defines to which group each element belongbelongToGroup = np.array([0, 1, 0, 2])nGroups=np.max(belongToGroup);# Calculate a matrix obtained by summing elements belonging to the same groupM_sum = np.zeros((nGroups+1,nGroups+1))for g1 in range(nGroups+1):    idxG1 = np.where(belongToGroup==g1)    for g2 in range(nGroups+1):        idxG2 = np.where(belongToGroup==g2)        print('g1 = ' + str(g1))        print('g2 = ' + str(g2))        print(idxG1[0])        print(idxG2[0])        print(M[idxG1[0],idxG2[0]])        print(np.sum(M[idxG1[0],idxG2[0]]))        M_sum[g1,g2]=np.sum(M[idxG1[0],idxG2[0]])print('')print('Example of the problem:')print('Elements I would like to sum to obtain M_sum[0,0]')print(M[0:2,0:2])print('Elements that are summed instead')print(M[[0,1],[0,1]])问题示例：在上面的例子中，元素 M_sum[0,0] 应该是 M[0,0]、M[0,1]、M[1,0] 和 M[1,1] 的和] 相反，它被计算为 M[0,0] 和 M[1,1] 的总和

Answer 1

在 MATLAB 中，使用 2 个列表（实际上是矩阵）进行索引会选择一个块。 numpy另一方面，尝试相互广播索引数组，并返回选定的点。它的行为sub2ind与 MATLAB 中的行为接近。

In [971]: arr = np.arange(16).reshape(4,4)                                      

In [972]: arr                                                                   

Out[972]: 

array([[ 0,  1,  2,  3],

       [ 4,  5,  6,  7],

       [ 8,  9, 10, 11],

       [12, 13, 14, 15]])

In [973]: i1, i2 = np.array([0,2,3]), np.array([1,2,0])                         

使用 2 个相同大小的一维数组进行索引：

In [974]: arr[i1,i2]

Out[974]: array([ 1, 10, 12])

这实际上[arr[0,1], arr[2,2], arr[3,0]]为匹配索引的每个点返回一个元素。

但是如果我把一个索引变成一个“列向量”，它会从行中i2选择，而从列中选择。

In [975]: arr[i1[:,None], i2]                                                   

Out[975]: 

array([[ 1,  2,  0],

       [ 9, 10,  8],

       [13, 14, 12]])

MATLAB 使块索引变得容易，而单独访问则更难。numpy尽管底层机制是相同的，但在块中访问有点困难。

使用您的示例i1[0]，i2[0]可以是如下数组：

array([0, 2]), array([3])

(2,) (1,)

形状 (1,) 数组也可以与 (2,) 或 (2,1) 数组一起广播。如果代替代码将失败is[0]是np.array([0,1,2])，A（3）阵列，其不能与（2）阵列对。但是对于 (2,1) 它会产生一个 (2,3) 块。

Answer 2

您可以使用np.ix_获取 matlab-ish 行为：

A = np.arange(9).reshape(3, 3)

A[[1,2],[0,2]]

# array([3, 8])

A[np.ix_([1,2],[0,2])]

# array([[3, 5],

#        [6, 8]])

在np.ix_幕后，@hpaulj 详细描述了以下内容：

np.ix_([1,2],[0,2])

# (array([[1],

#        [2]]), array([[0, 2]]))

您可以将其应用于您的特定问题，如下所示：

M = np.random.randint(0, 10, (n, n))

M

# array([[6, 2, 7, 1],

#        [6, 7, 9, 5],

#        [9, 4, 3, 2],

#        [3, 1, 7, 9]])

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

ng = idx.max() + 1

out = np.zeros((ng, ng), M.dtype)

np.add.at(out, np.ix_(idx, idx), M)

out

# array([[25,  6,  3],

#        [15,  7,  5],

#        [10,  1,  9]])

顺便说一句：有一个更快但不太明显的解决方案，它依赖于平面索引：

np.bincount(np.ravel_multi_index(np.ix_(idx, idx), (ng, ng)).ravel(), M.ravel(), ng*ng).reshape(ng, ng)

# array([[25.,  6.,  3.],

#        [15.,  7.,  5.],

#        [10.,  1.,  9.]])