4

In this example I'm trying to create a vector by selecting relevant elements from a multidimensional array.

#data
n=3
rng = 4
x = np.array([0,1,2],dtype=int)
y = np.array([0,3,1],dtype=int)
P = np.reshape(np.arange(n*rng*rng),(n,rng,rng))

output = np.zeros(n)
for i in range(n):
    output[i] = P[i,x[i],y[i]] 

This returns

array([  0.,  23.,  41.]) 

Now I'm trying to vectorize the above operation. To me, the logical thing would be to set

output = P[0:n,x,y]

but this returns

array([[ 0,  7,  9],
       [16, 23, 25],
       [32, 39, 41]])

Can anybody explain what is going on here and what I should do to obtain the intended output?

Thanks in advance

1 Answer 1

3

All you need is:

>>> P[np.arange(n), x, y]
array([ 0, 23, 41])

Related: Indexing Multi-dimensional arrays

Sign up to request clarification or add additional context in comments.

1 Comment

Thanks a lot! Been struggling with this all day.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.