Get indices for all elements in an array in numpy

Question

I'm trying to get a list of the indices for all the elements in an array so for an array of 1000 x 1000 I end up with [(0,0), (0,1),...,(999,999)].

I made a function to do this which is below:

def indices(alist):
    results = []
    ele = alist.size
    counterx = 0
    countery = 0
    x = alist.shape[0]
    y = alist.shape[1]
    while counterx < x:
        while countery < y:
            results.append((counterx,countery))
            countery += 1
        counterx += 1
        countery = 0
    return results

After I timed it, it seemed quite slow as it was taking about 650 ms to run (granted on a slow laptop). So, figuring that numpy must have a way to do this faster than my mediocre coding, I took a look at the documentation and tried:

indices = [k for k in numpy.ndindex(q.shape)]
which took about 4.5 SECONDS (wtf?)
indices = [x for x,i in numpy.ndenumerate(q)]
better, but 1.5 seconds!

Is there a faster way to do this?

Thanks

score 11 · Answer 1 · answered Jun 28 '13 at 20:14

11

how about np.ndindex?

np.ndindex(1000,1000)

This returns an iterable object:

>>> ix = numpy.ndindex(1000,1000)
>>> next(ix)
(0, 0)
>>> next(ix)
(0, 1)
>>> next(ix)
(0, 2)

In general, if you have an array, you can build the index iterable via:

index_iterable = np.ndindex(*arr.shape)

Of course, there's always np.ndenumerate as well which could be implemented like this:

def ndenumerate(arr):
    for ix in np.ndindex(*arr.shape):
        yield ix,arr[ix]

answered Jun 28 '13 at 20:14

mgilson

300,191
65
633
696

Out of idle curiosity, how do the timings for these compare to the timings made by the OP? – tripleee Jun 29 '13 at 09:02

john_science · Answer 2 · 2013-06-28T20:21:31.857

4

Have you thought of using itertools? It will generate an iterator for your results, and will almost certainly be optimally fast:

import itertools

a = range(1000)
b = range(1000)

product = itertools.product(a, b)

for x in product:
    print x

# (0, 0)
# (0, 1)
# ...
# (999, 999)

Notice this didn't require the dependency on numpy. Also, notice the fun use of range to create a list from 0 to 999.

edited Jun 28 '13 at 20:21

answered Jun 28 '13 at 20:12

john_science

6,325
6
43
60

good answer assuming it is a true 2d array and not a list of variably sized list (+1) – Joran Beasley Jun 28 '13 at 20:13
Ah, hmm. Yes. Well, I answered the question the OP asked. I can expand if the question is more general. – john_science Jun 28 '13 at 20:14

score 3 · Accepted Answer · edited May 23 '17 at 12:32

3

Ahha!

Using numpy to build an array of all combinations of two arrays

Runs in 41 ms as opposed to the 330ms using the itertool.product one!

edited May 23 '17 at 12:32

Community

1
1

answered Jun 29 '13 at 20:51

Jason White

666
4
10
23

Get indices for all elements in an array in numpy

3 Answers3