Using elements of numpy array as index of another numpy array

Viewed 96

I have an unknown 2D numpy array like the following:

a = np.array([
    [1, 2],
    [0, 2],
])

I want to use the elements of the array as an index to another array b to produce the following effect:

b = np.random.randn(5, 5)
c = b[[1, 2], [0, 2]]

How can I use the variable a to replace the hardcoded values in the index?

Using the following code did not work:

b = np.random.randn(5, 5)
c = [*a]

As the * expression is used in an index.

2 Answers

Just change the outer layer of a to a tuple to trigger advanced indexing: b[tuple(a)].

This is pretty much how the index syntax of b[[1, 2], [0, 2]] is interpreted; the elements are changed to NumPy arrays and are put into a tuple.

I don't know if I got it but, try using a list for the index instead of an np array:

import numpy as np

a = [
    [1, 2],
    [0, 2]
    ]

b = np.random.randn(5, 5)
c = b[a]
Related