在Python中汇总矩阵的选定元素

时间:2019-04-02 17:19:37

标签: python arrays numpy matrix indices

我有一个[n x n]矩阵,其中包含属于不同组的值;还有一个[1 x n]向量,定义每个元素属于哪个组。 (n通常为〜1E4,在此示例中为n = 4)

我想计算一个矩阵,该矩阵将属于同一组的所有元素加在一起。

我使用np.where()来计算每个组的元素所在的索引。 当我使用计算的索引时,我没有获得预期的元素,因为我选择了成对的位置而不是范围(我习惯于Matlab,在这里我可以简单地选择M(idx1,idx2))。

import numpy as np

n=4
M = np.random.rand(n,n)
print(M)

# This vector defines to which group each element belong
belongToGroup = np.array([0, 1, 0, 2])

nGroups=np.max(belongToGroup);

# Calculate a matrix obtained by summing elements belonging to the same group
M_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]

2 个答案:

答案 0 :(得分:2)

在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个相同大小的1d数组建立索引:

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])(一个无法与(2,)数组配对的(3,)数组,则您的代码将失败。但是有了(2,1),它就会产生(2,3)块。

答案 1 :(得分:1)

您可以使用np.ix_获得类似matlab的行为:

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.]])