CODEWITHSUNDEEP
pythonmultidimensional-arraynumpy
Aso Agile
Aso Agile

Reputation: 417

Numpy array operation using another indices array

I want to do a multidimensional array operation using numpy on three arrays, of which one is an index array, e.g.:

a = numpy.arange(20).reshape((5, 4))
# a = [[ 0  1  2  3]  [ 4  5  6  7]  [ 8  9 10 11]  [12 13 14 15]  [16 17 18 19]]

b = numpy.arange(24).reshape(((3, 2, 4)))
# b = [[[ 0  1  2  3]   [ 4  5  6  7]]   [[ 8  9 10 11]   [12 13 14 15]]  [[16 17 18 19]    [20 21 22 23]]]

c = numpy.array([0,0,1,1,2])
# c = [0 0 1 1 2]

now, what I want is:

d = a * b[&] + b[&&]

where & is the second element of second dimension of b (e.g. [ 4 5 6 7]) and && is the first element of second dimension (e.g. [ 0 1 2 3]) related to i-th item of the first dimension of b, where i is from array c (e.g. c[0]=0 for the first element of first dimension of array b). d has same dimension as a.

Edit: Answer for the above example is: # d = [[0 6 14 24] [16 26 38 52] [104 126 150 176] [152 178 206 236] [336 374 414 456]]

Thanks

Upvotes: 1

Views: 221

Answers (1)

eumiro
eumiro

Reputation: 212885

>>> a * b[c,1,:] + b[c,0,:]
array([[  0,   6,  14,  24],
       [ 16,  26,  38,  52],
       [104, 126, 150, 176],
       [152, 178, 206, 236],
       [336, 374, 414, 456]])

Upvotes: 3

Related Questions