Generate NumPy array containing the indices of another NumPy array

Question

I'd like to generate a np.ndarray NumPy array for a given shape of another NumPy array. The former array should contain the corresponding indices for each cell of the latter array.

Example 1

Let's say we have a = np.ones((3,)) which has a shape of (3,). I'd expect

[[0]
 [1]
 [2]]

since there is a[0], a[1] and a[2] in a which can be accessed by their indices 0, 1 and 2.

Example 2

For a shape of (3, 2) like b = np.ones((3, 2)) there is already very much to write. I'd expect

[[[0 0]
  [0 1]]

 [[1 0]
  [1 1]]

 [[2 0]
  [2 1]]]

since there are 6 cells in b which can be accessed by the corresponding indices b[0][0], b[0][1] for the first row, b[1][0], b[1][1] for the second row and b[2][0], b[2][1] for the third row. Therefore we get [0 0], [0 1], [1 0], [1 1], [2 0] and [2 1] at the matching positions in the generated array.

Thank you very much for taking the time. Let me know if I can clarify the question in any way.

Your first array has shape (3, 1), and the second one shape (3, 2, 2). How do you go from (3,) to (3, 1) and from (3, 2) to (3, 2, 2)? Moreover, what determines the contents of the resulting arrays? — fuglede
– fuglede, Commented Jul 11, 2018 at 21:14
What would you expect if the inputs were, say, (4, 1), or (4, 3, 2)? — fuglede
– fuglede, Commented Jul 11, 2018 at 21:29
@fuglede Thank you for your comment. I tried to clarify the question a little bit further. Regarding your question in the comment: (3,) results in (3, 1) since there are 3 cells which can be accessed by 1 index - (3, 2) results in (3, 2, 2) since there are 6 cells shaped 3 rows and 2 columns with each 2 indices to uniquely access them. So at first the generated array has the same shape as the input array - plus one additional dimension for the number of indices. (4, 1) would therefore result in (4, 1, 2) and (4, 3, 2) would result in (4, 3, 2, 3). I hope this helps? :-) — finefoot
– finefoot, Commented Jul 12, 2018 at 10:27
Yep, that's clearer, thanks. In that case, it would seem that the provided answer is what you're looking for? — fuglede
– fuglede, Commented Jul 12, 2018 at 10:56
@fuglede Yes, I think so, too. Just wanted to clear that up in case you'd have another idea or answer. :-) — finefoot
– finefoot, Commented Jul 12, 2018 at 10:58

akuiper · Accepted Answer · 2018-07-11 21:47:13Z

4

One way to do it with np.indices and np.stack:

np.stack(np.indices((3,)), -1)

#array([[0],
#       [1],
#       [2]])

np.stack(np.indices((3,2)), -1)

#array([[[0, 0],
#        [0, 1]],
#       [[1, 0],
#        [1, 1]],
#       [[2, 0],
#        [2, 1]]])

np.indices returns an array of index grid where each subarray represents an axis:

np.indices((3, 2))

#array([[[0, 0],
#        [1, 1],
#        [2, 2]],        
#       [[0, 1],
#        [0, 1],
#        [0, 1]]])

Then transpose the array with np.stack, stacking index for each element from different axis:

np.stack(np.indices((3,2)), -1)

#array([[[0, 0],
#        [0, 1]],
#       [[1, 0],
#        [1, 1]],
#       [[2, 0],
#        [2, 1]]])

edited Jul 11, 2018 at 21:47

answered Jul 11, 2018 at 21:33

akuiper

216k33 gold badges362 silver badges379 bronze badges

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

Collectives™ on Stack Overflow

Generate NumPy array containing the indices of another NumPy array

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related