基于索引数组的numpy数组的和

djmepvbi  于 2023-04-30  发布在  其他
关注(0)|答案(2)|浏览(128)

在Python 3中,使用numpy,我有一个值数组和索引数组。我需要找到值数组的总和(或平均值)索引数组中的每个索引,索引数组中的每个值。出于速度的原因,我希望使用numpy切片而不是for循环来实现这一点。

import numpy
value_array = numpy.array([[0.1, 0.2, 0.5, 1 ],)
                           [2  ,   5,  10, 20]])
index_array = numpy.array([[  0,   1,   1, 2 ],
                           [  2,   0,   1, 4 ]])
# Something like desired_sum = numpy.sum(value_array[index_array])
# Where the output will be
desired_sum  = numpy.array([5.1, 10.7, 3, 0, 20])

这将在具有几百个索引的较大数组{~shape =(2000,2000)}上运行。

pkwftd7m

pkwftd7m1#

This is sort of a duplicate of this question. Python 3的代码是:

import numpy
value_array = numpy.array([[0.1, 0.2, 0.5, 1 ],
                           [2  ,   5,  10, 20]])
# index_array must be integers an have no negative values
index_array = numpy.array([[  0,   1,   1, 2 ],
                           [  2,   0,   1, 4 ]]).astype(int)
sum_array     = numpy.bincount(index_array.flatten(), weights = value_array.flatten())
average_array = sum_array / numpy.bincount(index_array.flatten())

print('sum_array = ' + repr(sum_array))
print('average_array = ' + repr(average_array))

输出为

sum_array = array([ 5.1, 10.7, 3. , 0. , 20. ])
average_array = array([ 2.55, 3.56666667, 1.5, nan, 20. ])
bttbmeg0

bttbmeg02#

您可以在结果位置的范围内广播索引和值矩阵,再添加一个维度以形成值矩阵的多个掩码。
例如,值将给予2d矩阵的5个副本:

array([[[ 0.1,  0.2,  0.5,  1. ],
        [ 2. ,  5. , 10. , 20. ]],

       [[ 0.1,  0.2,  0.5,  1. ],
        [ 2. ,  5. , 10. , 20. ]],

       [[ 0.1,  0.2,  0.5,  1. ],
        [ 2. ,  5. , 10. , 20. ]],

       [[ 0.1,  0.2,  0.5,  1. ],
        [ 2. ,  5. , 10. , 20. ]],

       [[ 0.1,  0.2,  0.5,  1. ],
        [ 2. ,  5. , 10. , 20. ]]])

对于这些副本中的每一个,仅保留与相应index_array中的索引相对应的值(使用0..4输出指标)。鉴于numpy只处理固定大小的维度,选择/排除值将通过乘以1或0(索引比较的布尔值)来完成:

import numpy
value_array = numpy.array([[0.1, 0.2, 0.5, 1 ],
                           [2  ,   5,  10, 20]])
index_array = numpy.array([[  0,   1,   1, 2 ],
                           [  2,   0,   1, 4 ]])

result_index = numpy.arange(5)[:,None,None]

eligible_values = numpy.tile(value_array,(5,1)).reshape(-1,*value_array.shape)
eligible_values *= result_index==index_array
result = numpy.sum(eligible_values,axis=(1,2))

print(result)
[ 5.1 10.7  3.   0.  20. ]

要获得平均值,您可以以类似的方式计算计数(您只需要使用这些计数的索引):

counts = numpy.sum(result_index==index_array,axis=(1,2))

print(counts)
[2 3 2 0 1]

print(result/counts)  # averages
[ 2.55        3.56666667  1.5                nan 20.        ]

相关问题