logo标签列表

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

pythonnumpyindices
6

我正在尝试获取数组中所有元素的索引列表，因此对于一个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毫秒（在一台较慢的笔记本电脑上，操作系统是Windows 10）。所以，我想到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!

有没有更快的方法来完成这个任务？

谢谢。

3个回答
11

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]
闲来无事，这些时间与 OP 的时间相比如何？ - tripleee
4

你是否考虑过使用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 的列表。

好的回答，假设它是一个真正的二维数组，而不是一个大小可变的列表（+1）。 - Joran Beasley
啊，嗯。是的。好吧，我回答了OP提出的问题。如果问题更普遍，我可以进行扩展。 - john_science
3
网页内容由stack overflow 提供, 点击上面的
可以查看英文原文，
原文链接