Is there an easy way to index a numpy multidimensional array along the last dimension, using an array of indices? For example, take an array a
of shape (10, 10, 20)
. Let's assume I have an array of indices b
, of shape (10, 10)
so that the result would be c[i, j] = a[i, j, b[i, j]]
.
I've tried the following example:
a = np.ones((10, 10, 20))
b = np.tile(np.arange(10) + 10, (10, 1))
c = a[b]
However, this doesn't work because it then tries to index like a[b[i, j], b[i, j]]
, which is not the same as a[i, j, b[i, j]]
. And so on. Is there an easy way to do this without resorting to a loop?
There are several ways to do this. Let's first generate some test data:
In [1]: a = np.random.rand(10, 10, 20)
In [2]: b = np.random.randint(20, size=(10,10)) # random integers in range 0..19
One way to solve the question would be to create two index vectors, where one is a row vector and the other a column vector of 0..9 using meshgrid:
In [3]: i1, i0 = np.meshgrid(range(10), range(10), sparse=True)
In [4]: c = a[i0, i1, b]
This works because i0
, i1
and b
will all be broadcasted to 10x10 matrices. Quick test for correctness:
In [5]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10))
Out[5]: True
Another way would be to use choose and rollaxis:
# choose needs a sequence of length 20, so move last axis to front
In [22]: aa = np.rollaxis(a, -1)
In [23]: c = np.choose(b, aa)
In [24]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10))
Out[24]: True
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With