32

For python dict, I could use iteritems() to loop through key and value at the same time. But I cannot find such functionality for NumPy array. I have to manually track idx like this:

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

Is there a better way to do this?

jpp
  • 159,742
  • 34
  • 281
  • 339
user40780
  • 1,828
  • 7
  • 29
  • 50
  • 1
    You can use `enumerate` for this.. https://stackoverflow.com/questions/22171558/what-does-enumerate-mean. Or iterate `range(len(theta))`. – jpp Mar 20 '18 at 12:51
  • 2
    Or better, `ndenumerate` if you are dealing with NDarrays. – cs95 Dec 28 '18 at 20:35

2 Answers2

61

There are a few alternatives. The below assumes you are iterating over a 1d NumPy array.

Iterate with range

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

Note this is the only of the 3 solutions which will work with numba. This is noteworthy since iterating over a NumPy array explicitly is usually only efficient when combined with numba or another means of pre-compilation.

Iterate with enumerate

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

The most efficient of the 3 solutions for 1d arrays. See benchmarking below.

Iterate with np.ndenumerate

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

Notice the additional indexing step in idx[0]. This is necessary since the index (like shape) of a 1d NumPy array is given as a singleton tuple. For a 1d array, np.ndenumerate is inefficient; its benefits only show for multi-dimensional arrays.

Performance benchmarking

# 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
jpp
  • 159,742
  • 34
  • 281
  • 339
10

You can use numpy.ndenumerate for example

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

For more information refer to https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndenumerate.html

Frank Bryce
  • 8,076
  • 4
  • 38
  • 56
sushmit
  • 4,369
  • 2
  • 35
  • 38
  • 2
    Also keep in mind that this is extensible to multiple dimensions. index is a tuple containing the index of each dimension of the array. Quite useful! – Frank Bryce Apr 10 '19 at 18:39