如何使用索引和值迭代1d NumPy数组

Question

对于python dict,我可以使用iteritems()同时循环键和值.但我找不到NumPy阵列的这种功能.我必须idx像这样手动跟踪:

idx = 0 
for j in theta:
   some_function(idx,j,theta)
   idx += 1

有一个更好的方法吗？

Answer 1

还有一些选择.以下假设您正在迭代1d NumPy数组.

迭代 range

for j in range(theta.shape[0]):  # or range(len(theta))
   some_function(j, theta[j], theta)

请注意,这是可以使用的3种解决方案中唯一的解决方案numba.这是值得注意的,因为显式地迭代NumPy数组通常仅在与numba预编译或其他预编译方式结合时才有效.

迭代 enumerate

for idx, j in enumerate(theta):
   some_function(idx, j, theta)

1d阵列的3种解决方案中效率最高的.请参阅下面的基准测试

迭代 np.ndenumerate

for idx, j in np.ndenumerate(theta):
   some_function(idx[0], j, theta)

请注意其他索引步骤idx[0].这是必要的,因为shape1d NumPy数组的索引(如)作为单个元组给出.对于1d阵列,np.ndenumerate效率低; 它的好处只适用于多维数组.

绩效基准

# Python 3.7, NumPy 1.14.3

np.random.seed(0)

arr = np.random.random(10**6)

def enumerater(arr):
    for index, value in enumerate(arr):
        index, value
        pass

def ranger(arr):
    for index in range(len(arr)):
        index, arr[index]
        pass

def ndenumerater(arr):
    for index, value in np.ndenumerate(arr):
        index[0], value
        pass

%timeit enumerater(arr)    # 131 ms
%timeit ranger(arr)        # 171 ms
%timeit ndenumerater(arr)  # 579 ms

Answer 2

您可以使用numpy.ndenumerate例如

import numpy as np
test_array = np.arange(2, 3, 0.1)
for index, value in np.ndenumerate(test_array):
    print(index[0], value)

有关更多信息，请参阅https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndenumerate.html