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]

Answer 2

您是否考虑过使用itertools？它将为您的结果生成一个迭代器，几乎肯定会以最佳的速度运行：

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)

请注意，这不需要依赖numpy。另外，请注意使用range来创建0到999的列表。

Answer 3

AHHA！

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

运行时间为41毫秒，而使用itertool.product的运行速度为330毫秒！

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

3 个答案: