martianwars martianwars -4 years ago 75
Python Question

Vectorized way to add elements using an index map?

I have two

numpy
arrays, one of size
(386, 3, 4)
and another of size
(386, 4)
, which I will call
values
and
keys
respectively. The second array contains integers which are indices to my output array. I need to implement the following
for
loop -

for i in range(386):
for j in range(4):
output[keys[i, j]] += values[i, :, j]


Of course,
output
has dimensions
(max_index + 1, 3)
. Could I make way with a vectorized implementation?

Answer Source

I think np.add.at should do what you want:

np.add.at(output, keys, np.transpose(values, (0, 2, 1)))

Small array example:

values
# array([[[100, 200, 300, 400],
    [ 10,  20,  30,  40],
    [  1,   2,   3,   4]],

   [[500, 600, 700, 800],
    [ 50,  60,  70,  80],
    [  5,   6,   7,   8]]])
keys
# array([[4, 0, 3, 1],
   [1, 0, 2, 2]])
out
# array([[0, 0, 0],
   [0, 0, 0],
   [0, 0, 0],
   [0, 0, 0],
   [0, 0, 0]])
np.add.at(out, keys, np.transpose(values, (0, 2, 1)))
out
# array([[ 800,   80,    8],
   [ 900,   90,    9],
   [1500,  150,   15],
   [ 300,   30,    3],
   [ 100,   10,    1]])
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download