Get array of indices for array

Question

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?

Another method that just occurred to me is np.transpose(np.nonzero(a == a)), but that's almost more weird than using indices and reshape — squirl
– squirl, Commented Mar 1, 2018 at 10:39

Kasravnd · Accepted Answer · 2018-03-01 11:42:25Z

3

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)

edited Mar 1, 2018 at 11:42

answered Mar 1, 2018 at 10:46

Kasravnd

108k19 gold badges167 silver badges195 bronze badges

Sign up to request clarification or add additional context in comments.

6 Comments

squirl Over a year ago

The order of the indices doesn't really matter, I guess. I'll go with the first one. Thanks!

Paul Panzer Over a year ago

The first one doesn't really work. Try it on a = np.zeros((2, 2)).

Paul Panzer Over a year ago

Actually, already in the (3,3) example: you can see (1, 2) occurring twice, while (2, 1) is missing

squirl Over a year ago

Hmm.. good point. Perhaps I'll use np.transpose(np.indices(a.shape)).reshape(a.size, 2) instead

Kasravnd Over a year ago

@Samadi Noe just add a T in the middle of two methods.

|

Collectives™ on Stack Overflow

Get array of indices for array

1 Answer 1

6 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

6 Comments

Your Answer

Sign up or log in

Post as a guest

Related