How to index an array with its indices in numpy?

Question

For example, I have an array and its indices:

  xx = np.arange(6).reshape(2,3)
  index = np.argwhere(xx>3)

If I use

 xx[index]

I get an error:

IndexError: index 2 is out of bounds for axis 0 with size 2

But if I write:

xx[index[0], index[1]]

I get:

array([4, 5])

Which is exactly what I want. But if I have an N dimensional array, do I need to write

xx[index[0], index[1], ..., index[N]]

in this case, or is there a better way to do it?

Answer 1

From the documentation of numpy:

The output of argwhere is not suitable for indexing arrays. For this purpose use nonzero(a) instead.

So, this would be the recommended way and it works just fine:

>>> xx = np.arange(6).reshape(2,3)
>>> index = np.nonzero(xx>3)
>>> xx[index]
array([4, 5])

---

The point of np.argwhere is essentially that it will group the indices by element and so the resulting output will be of the shape (M, xx.ndim) where M is the number of non-zero elements, which makes it clearly not suitable for indexing.

np.nonzero, on the other hand, groups indexes by dimension, returning for each dimension a tuple which is what you need!

So, basically if the elements at [x1,y1], [x2,y2] and [x3, y3] meet your condition, argwhere will return an array like this:

[[x1, y1], [x2, y2], [x3, y3]]

while nonzero will return a tuple:

([x1, x2, x3], [y1, y2, y3])

Note that the latter is what you need for indexing. Also, the latter is also almost the transpose of the former, so if you are stuck for some reason with an array of the form produced by argwhere you can transpose and convert it to a tuple and you are golden. So tuple(np.transpose(index)) in your case. If you are using argwhere to get your indices then I'd recommend to simply use nonzero instead.