NumPy - Iterating over Arrays

NumPy contains an iterator object numpy.nditer, which is an efficient multi-dimensional iterator object to iterate over arrays. Each element is provided one by one using the standard Python iterator interface.

In the example below, where an array is created using arange() function and then nditer is used to iterate over it.

import numpy as np

#creating an array
Arr = np.arange(10,25)
Arr = Arr.reshape(3,5)

#printing the array
print("The original array is:")
print(Arr)

#iterating over the array using nditer
print("\nThe modified array is:")
for i in np.nditer(Arr):
  print(i, end=" ")

The output of the above code will be:

The original array is:
[[10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]

The modified array is:
10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

The order of iteration is chosen to match the memory layout of an array instead of using a standard C or Fortran ordering. This is done for access efficiency, reflecting the idea that by default one simply wants to visit each element without concern for a particular ordering. This can be seen by iterating over the transpose of the above array.

import numpy as np

#creating an array
Arr = np.arange(10,25)
Arr = Arr.reshape(3,5)

#printing the array
print("The original array is:")
print(Arr)

#printing the array
print("\nThe transpose of the array is:")
print(Arr.T)

#iterating over the transpose of the array
print("\nThe modified array is:")
for i in np.nditer(Arr.T):
  print(i, end=" ")

The output of the above code will be:

The original array is:
[[10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]

The transpose of the array is:
[[10 15 20]
 [11 16 21]
 [12 17 22]
 [13 18 23]
 [14 19 24]]

The modified array is:
10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

Iteration order

There are instances when it is important to visit the elements of an array in a specific order, irrespective of the layout of the elements in memory. The nditer object provides an order parameter to control this aspect of iteration. The default is order='K' to keep the existing order of the array. This can be specified as order='C' for C order or order='F' for Fortran order.

import numpy as np

#creating an array
Arr = np.arange(10,25)
Arr = Arr.reshape(3,5)

#printing the array
print("The original array is:")
print(Arr)

#iterating over the array using order='K'
print("\nThe modified array (order='K'):")
for i in np.nditer(Arr, order='K'):
  print(i, end=" ")

#iterating over the array using order='C'
print("\n\nThe modified array (order='C'):")
for i in np.nditer(Arr, order='C'):
  print(i, end=" ")

#iterating over the array using order='F'
print("\n\nThe modified array (order='F'):")
for i in np.nditer(Arr, order='F'):
  print(i, end=" ")

The output of the above code will be:

The original array is:
[[10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]

The modified array (order='K'):
10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 

The modified array (order='C'):
10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 

The modified array (order='F'):
10 15 20 11 16 21 12 17 22 13 18 23 14 19 24

Modifying Array Values

The nditer object has another optional parameter called op_flags. By default, the nditer treats the input operand as a read-only object. To be able to modify the array elements, it must be specified either read-write or write-only mode using op_flags=['readwrite'] or op_flags=['writeonly'].

import numpy as np

#creating an array
Arr = np.arange(10,25)
Arr = Arr.reshape(3,5)

#printing the array
print("The original array is:")
print(Arr)

#modifying the elements of the array
for x in np.nditer(Arr, op_flags=['readwrite']):
  x[...] = 2*x

#printing the array
print("\nThe modified array is:")
print(Arr)

The output of the above code will be:

The original array is:
[[10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]

The modified array is:
[[20 22 24 26 28]
 [30 32 34 36 38]
 [40 42 44 46 48]]

External Loop

The nditer object has another optional parameter called flags, which can take the following values:

Parameter	Description
c_index	C-order index to be tracked.
f_index	Fortran-order index to be tracked.
multi_index	Multi-index, or a tuple of indices with one per iteration dimension, to be tracked.
external_loop	Causes values given to be one-dimensional arrays with multiple values instead of zero-dimensional arrays.
ranged	Allows the iterator to be restricted to a sub-range of the iterindex values.

In the example below, one-dimensional arrays corresponding to each column is traversed by the iterator.

import numpy as np

#creating an array
Arr = np.arange(10,25)
Arr = Arr.reshape(3,5)

#printing the array
print("The original array is:")
print(Arr)

#iterating over the array using external_loop
print("\nThe modified array is:")
for i in np.nditer(Arr, flags = ['external_loop'], order = 'F'):
  print(i, end=" ")

The output of the above code will be:

The original array is:
[[10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]

The modified array is:
[10 15 20] [11 16 21] [12 17 22] [13 18 23] [14 19 24]

Broadcasting Array Iteration

If two arrays are broadcastable, a combined nditer object is able to iterate upon them concurrently.

In the example below, the nditer object is used to iterate over arrays arr1 and arr2 concurrently. The dimension of arr1 is (2x3) and dimension of arr2 is (1x3).

import numpy as np

#creating first array
arr1 = np.arange(10,70,10)
arr1 = arr1.reshape(2,3)

#creating second array
arr2 = np.array([1, 2, 3], dtype = int)

#printing arrays
print("arr1 is:")
print(arr1)
print("\narr2 is:")
print(arr2)

#iterating over two arrays concurrently
print("\nThe modified array is:")
for x,y in np.nditer([arr1, arr2]):
  print("%d:%d" % (x,y), end=" ")

The output of the above code will be:

arr1 is:
[[10 20 30]
 [40 50 60]]

arr2 is:
[1 2 3]

The modified array is:
10:1 20:2 30:3 40:1 50:2 60:3