Gets the index array of the array… here is a solution to the problem.
Gets the index array of the array
If I have such a multidimensional array:
a = np.array([[9,9,9],[9,0,9],[9,9,9]])
I want to get the array for each index in that array as follows:
i = np.array([[0,0],[0,1],[0,2],[1,0],[1,1],...])
One approach I’ve found is this, using np.indices
:
i = np.transpose(np.indices(a.shape)).reshape(a.shape[0] * a.shape[1], 2)
But this seems a bit clumsy, especially considering the existence of np.nonzero
, which pretty much satisfies my request.
Is there a built-in numpy function that generates an indexed array of each item in a 2D numpy array?
Solution
Here’s a more concise way (if the order isn’t important):
In [56]: np.indices(a.shape). T.reshape(a.size, 2)
Out[56]:
array([[0, 0],
[1, 0],
[2, 0],
[0, 1],
[1, 1],
[2, 1],
[0, 2],
[1, 2],
[2, 2]])
If you want to use it in the expected order, you can use dstack
:
In [46]: np.dstack(np.indices(a.shape)).reshape(a.size, 2)
Out[46]:
array([[0, 0],
[0, 1],
[0, 2],
[1, 0],
[1, 1],
[1, 2],
[2, 0],
[2, 1],
[2, 2]])
For the first method,
if you don’t want to use reshape
, the other method is to use np.concatenate()
to concatenate along the first axis.
np.concatenate(np.indices(a.shape). T)