Home > OS >  How do I access the values of a numpy array with an array of indices?
How do I access the values of a numpy array with an array of indices?

Time:01-28

For example,

k = np.array([[[1,2,3,4],[1,2,3,4]]])
index = np.array([[0,0], [0,1]])

I want to be able to get the values from k responding to [0,0] and [0,1].

How could I do this?

If I use a for loop through the array it works.

for y in range(1):
    for x in range(1):
        k[index[y,x]]

However, I would like to do this without using for loops.

CodePudding user response:

In [50]: k = np.array([[[1,2,3,4],[1,2,3,4]]])
    ...: index = np.array([[0,0], [0,1]])

In [51]: k
Out[51]: 
array([[[1, 2, 3, 4],
        [1, 2, 3, 4]]])
In [52]: k.shape
Out[52]: (1, 2, 4)

Note the shape - 3d, due to the 3 levels of []

In [53]: index
Out[53]: 
array([[0, 0],
       [0, 1]])

Because this array is symmetric, it doesn't matter whether we use the rows or the columns. For a more general case you'll need to be clearer.

In any case, we index each dimension of k with an array

Using columns of index, and working with the first 2 dimensions:

In [54]: k[index[:,0],index[:,1]]
Out[54]: 
array([[1, 2, 3, 4],
       [1, 2, 3, 4]])

Looks much like k except it is 2d.

Or applying a 0 to the first size 1 dimension:

In [55]: k[0,index[:,0],index[:,1]]
Out[55]: array([1, 2])

Read more at https://numpy.org/doc/stable/user/basics.indexing.html

  • Related