A right proper numpy way to create flat list of indices into an array

user_381179 0
Viewed 98

Okay, this is super basic, which unfortunately means that searching for it gives a bajillion hits that all do something different and / or more complex.

Consider this code:

shape = (10, 20)
indices = []
for i in range(shape[0]):
    for j in range(shape[1]):
        indices.append([i, j])

or alternatively indices = itertools.product(range(10), range(20)).

Now, I feel like there must be a simple numpy function that does the same? Something like

indices = np.indices_into_shape((10, 20))

Most of the index-generating functions I can find via search generate multiple arrays, like in meshgrid or ix_.

numpy python
2 Answers

You can stack meshgrids:

np.dstack(np.meshgrid(np.arange(10), np.arange(20), indexing='ij')).reshape(-1, 2)
user_29887300

One way would be

np.argwhere(np.broadcast_to(True,(3,4)))
# array([[0, 0],
#        [0, 1],
#        [0, 2],
#        [0, 3],
#        [1, 0],
#        [1, 1],
#        [1, 2],
#        [1, 3],
#        [2, 0],
#        [2, 1],
#        [2, 2],
#        [2, 3]])

another (similar to @MadPhysicist's)

np.c_[np.unravel_index(np.arange(3*4),(3,4))]
user_72073920
Related