CODEWITHSUNDEEP
pythonarraysnumpymultidimensional-array
squirl
squirl

Reputation: 1784

Get array of indices for array

If I have a multidimensional array like this:

a = np.array([[9,9,9],[9,0,9],[9,9,9]])

I'd like to get an array of each index in that array, like so:

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

One way of doing this that I've found is like this, using np.indices:

i = np.transpose(np.indices(a.shape)).reshape(a.shape[0] * a.shape[1], 2)

But that seems somewhat clumsy, especially given the presence of np.nonzero which almost does what I want.

Is there a built-in numpy function that will produce an array of the indices of every item in a 2D numpy array?

Upvotes: 3

Views: 128

Answers (1)

Kasravnd
Kasravnd

Reputation: 107287

Here is one more concise way (if the order is not 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 it in your intended 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 approach if you don't want to use reshape another way is concatenation along the first axis using np.concatenate().

np.concatenate(np.indices(a.shape).T)

Upvotes: 3

Related Questions