Search code examples
pythonnumpyindices

Get indices for all elements in an array in numpy


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


Solution

  • 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!