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

Answer 1

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]

Answer 2

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.