Question

我需要一种方法来遍历任何形状的numpy数组中的每个元素，并将其索引存储在列表中。

此代码生成一般形状的数组：

import numpy as np

# Generate random shape for the array.
sha = np.random.randint(1, 5, size=(np.random.randint(1, 10)))
# Random array.
a =np.random.rand(*sha)

我需要迭代a中的每个元素并将其索引存储在列表中。

我最接近的是扁平阵列：

for i, elem in enumerate(a.flatten()):
    print i, elem

允许我迭代 a中的每个元素，无论它的形状如何，但我松开了每个元素存储在数组中的索引。

Answer 1

您可以使用itertools创建所有维度的product（笛卡儿）

from itertools import product
for i in product(*[range(dim) for dim in a.shape]):
    print i, a[i]

Answer 2

您正在寻找numpy.unravel_index：

>>> np.dstack(np.unravel_index(np.arange(a.size), a.shape))
array([[[0, 0, 0, ..., 0, 0, 0],
        [0, 0, 0, ..., 0, 0, 1],
        [0, 0, 0, ..., 0, 0, 2],
        ...,
        [2, 2, 0, ..., 1, 1, 1],
        [2, 2, 0, ..., 1, 1, 2],
        [2, 2, 0, ..., 1, 1, 3]]])

Answer 3

ndindex为任何数组生成多维索引元组

事实上，我们可以生成dtype对象的nd数组，并在一步中填充这些索引的一些函数（ok，2步）：

shape=(3,4,2)
A = np.empty(shape, dtype=object)
for i in np.ndindex(3,4,2):
    A[i]=i   # or use A[i]=list(i)

生成A：

array([[[(0, 0, 0), (0, 0, 1)],
        [(0, 1, 0), (0, 1, 1)],
        [(0, 2, 0), (0, 2, 1)],
        [(0, 3, 0), (0, 3, 1)]],

       [[(1, 0, 0), (1, 0, 1)],
        [(1, 1, 0), (1, 1, 1)],
        [(1, 2, 0), (1, 2, 1)],
        [(1, 3, 0), (1, 3, 1)]],

       [[(2, 0, 0), (2, 0, 1)],
        [(2, 1, 0), (2, 1, 1)],
        [(2, 2, 0), (2, 2, 1)],
        [(2, 3, 0), (2, 3, 1)]]], dtype=object)

我用元组而不是列表填充它，因为显示更清晰。在

array([[[[0, 0, 0], [0, 0, 1]],
        [[0, 1, 0], [0, 1, 1]],
        [[0, 2, 0], [0, 2, 1]],
        ....

很难区分数组维度和列表。

使用dtype对象，各个元素可以是任何东西。 np.empty填充None。

相关的数据结构将是这些元组的平面列表。它仍然可以通过使用np.ravel_multi_index进行访问。

L1 = [i for i in np.ndindex(shape)]

[(0, 0, 0),
 (0, 0, 1),
 (0, 1, 0),
 ...
 (2, 2, 1),
 (2, 3, 0),
 (2, 3, 1)]

作为模拟深层嵌套列表访问：

 L1[np.ravel_multi_index(i,shape)]

一些访问时间比较：

In [137]: %%timeit
for i in np.ndindex(shape):
    x=A[i]
10000 loops, best of 3: 88.3 µs per loop

In [138]: %%timeit
for i in np.ndindex(shape):
    x=L1[np.ravel_multi_index(i,shape)]
1000 loops, best of 3: 227 µs per loop

因此，数组的多维索引更快。

In [140]: %%timeit
   .....: for i in L1:
   .....:     x=i
1000000 loops, best of 3: 774 ns per loop

In [143]: %%timeit
   .....: for i in A.flat:
    x=i
1000000 loops, best of 3: 1.44 µs per loop

但是对列表的直接迭代要快得多。

表示itertools.produce迭代器：

In [163]: %%timeit
   .....: for i in itertools.product(*[range(x) for x in shape]):
    x=A[i]
   .....: 
100000 loops, best of 3: 12.1 µs per loop

迭代存储索引的嵌套数组

3 个答案: