python - 获取numpy中数组中所有元素的索引

Question

我正在尝试获取数组中所有元素的索引列表，因此对于 1000 x 1000 的数组，我最终得到 [(0,0), (0,1),..., (999,999)].

我做了一个函数来做到这一点，如下所示:

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

在我计时之后，它看起来很慢，因为它需要大约 650 毫秒才能运行(在一台慢速笔记本电脑上是这样)。因此，我认为 numpy 一定有办法比我的普通编码更快地执行此操作，我查看了文档并尝试了:

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!

有没有更快的方法来做到这一点？

谢谢

Answer 1

np.ndindex 怎么样？

np.ndindex(1000,1000)

这将返回一个可迭代对象:

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

一般来说，如果你有一个数组，你可以通过以下方式构建可迭代的索引:

index_iterable = np.ndindex(*arr.shape)

当然，总有 np.ndenumerate 也可以这样实现:

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