使用索引(numpy等效的python枚举)迭代numpy

Question

我正在尝试创建一个函数来计算多维numpy数组中元素之间的晶格距离(水平和垂直步长的数量).为此,我需要在迭代数组时从每个元素的索引中检索实际数字.我想将这些值存储为可以通过距离公式运行的数字.

对于示例阵列A.

 A=np.array([[1,2,3],[4,5,6],[7,8,9]])

我想创建一个迭代遍历每个元素的循环,对于第一个元素1,它将检索a = 0,b = 0,因为1是A [0,0],然后a = 0,b = 1表示元素2因为它位于A [0,1],依此类推......

我设想的输出是数组中每个元素的两个数字(对应于该元素的两个索引值).所以在上面的例子中,我将分配的两个值是a和b.我只需要在循环中检索这两个数字(而不是单独保存为另一个数据对象).

任何关于如何做到这一点的想法将不胜感激!

Answer 1

您可以遍历数组中的值,numpy.ndenumerate以获取数组中值的索引.

使用上面的文档:

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
for index, values in np.ndenumerate(A):
    print(index, values)  # operate here

Answer 2

您可以使用它来完成此操作np.ndenumerate，但通常不需要迭代数组。

您可以简单地创建一个网格（或开放网格）来一次获取所有索引，然后您可以更快地处理它们（矢量化）。

例如

>>> x, y = np.mgrid[slice(A.shape[0]), slice(A.shape[1])]
>>> x
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2]])
>>> y
array([[0, 1, 2],
       [0, 1, 2],
       [0, 1, 2]])

这些可以像任何其他数组一样进行处理。因此，如果您的函数需要索引可以进行矢量化，那么您不应该执行手动循环！

例如，计算每个点到一个点的晶格距离(2, 3)：

>>> abs(x - 2) + abs(y - 3)
array([[5, 4, 3],
       [4, 3, 2],
       [3, 2, 1]])

对于距离来说，速度ogrid会更快。只需替换np.mgrid为np.ogrid：

>>> x, y = np.ogrid[slice(A.shape[0]), slice(A.shape[1])]
>>> np.hypot(x - 2, y - 3)  # cartesian distance this time! :-)
array([[ 3.60555128,  2.82842712,  2.23606798],
       [ 3.16227766,  2.23606798,  1.41421356],
       [ 3.        ,  2.        ,  1.        ]])