pythonarraysnumpyindexing

Python: How to get values of a multi-dimensional array at certain indexed positions?

If I have a multi-dimensional numpy array like:

a = np.array([[0, 1, 2], 
              [3, 4, 5], 
              [6, 7, 8]])

How can I get the values at certain index positions in one step? For example, if were to define pairs of indices like:

indices = [[0,0], [1,1], [2,2]]

I would like:

a[indices] = [0, 4, 8]

Note that this does work for one-dimensional arrays (Python: How to get values of an array at certain index positions?), but I cannot see how to get this to work in more than one dimension. I am using Python 3.7.

Solution

  • As in the one-dimensional answer you linked, you can do this elegantly in 2 dimensions with numpy:

    a = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
rows, columns = zip([0, 0], [1, 1], [2, 2])
print(a[rows, columns])

    The output of the print will be:

    array([0, 4, 8])