NumPy Array Iteration

The NumPy iterator object numpy.nditer provides a flexible way to access one or multiple array elements.

The basic task of the iterator is to access array elements.

Next, we will create a 2x3 array using the arange() function and iterate over it using nditer.

Example

import numpy as np

a = np.arange(6).reshape(2, 3)
print('The original array is:')
print(a)
print('\n')
print('Iterating over the elements:')
for x in np.nditer(a):
    print(x, end=", ")
print('\n')

Output:

The original array is:
[[0 1 2]
 [3 4 5]]

Iterating over the elements:
0, 1, 2, 3, 4, 5,

The above example does not use the standard C or Fortran order. The chosen order is consistent with the array's memory layout, which is done to improve access efficiency. By default, it is row-major order (C-order).

This reflects the default behavior of accessing each element without considering its specific order. We can see this by iterating over the transpose of the above array and comparing it with the C-order access of the transposed copy, as shown in the following example:

Example

import numpy as np

a = np.arange(6).reshape(2, 3)
for x in np.nditer(a.T):
    print(x, end=", ")
print('\n')

for x in np.nditer(a.T.copy(order='C')):
    print(x, end=", ")
print('\n')

Output:

0, 1, 2, 3, 4, 5, 

0, 3, 1, 4, 2, 5,

From the above example, it can be seen that the traversal order of a and a.T is the same, which means their storage order in memory is also the same. However, the traversal result of a.T.copy(order='C') is different because it has a different storage method, with the default being row-wise access.

Controlling the Traversal Order

• for x in np.nditer(a, order='F'): Fortran order, column-major;
• for x in np.nditer(a.T, order='C'): C order, row-major;

Example

import numpy as np

a = np.arange(0, 60, 5)
a = a.reshape(3, 4)
print('The original array is:')
print(a)
print('\n')
print('The transpose of the original array is:')
b = a.T
print(b)
print('\n')
print('Sorted in C-style order:')
c = b.copy(order='C')
print(c)
for x in np.nditer(c):
    print(x, end=", ")
print('\n')
print('Sorted in F-style order:')
c = b.copy(order='F')
print(c)
for x in np.nditer(c):
    print(x, end=", ")

Output:

The original array is:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

The transpose of the original array is:
[[ 0 20 40]
 [ 5 25 45]
 [10 30 50]
 [15 35 55]]

Sorted in C-style order:
[[ 0 20 40]
 [ 5 25 45]
 [10 30 50]
 [15 35 55]]
0, 20, 40, 5, 25, 45, 10, 30, 50, 15, 35, 55, 

Sorted in F-style order:
[[ 0 20 40]
 [ 5 25 45]
 [10 30 50]
 [15 35 55]]
0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55,

The traversal order can be explicitly set to force the nditer object to use a specific order:

Example

import numpy as np

a = np.arange(0, 60, 5)
a = a.reshape(3, 4)
print('The original array is:')
print(a)
print('\n')
print('Sorted in C-style order:')
for x in np.nditer(a, order='C'):
    print(x, end=", ")
print('\n')
print('Sorted in F-style order:')

for x in np.nditer(a, order='F'):
    print(x, end=", ")

Output result:

The original array is:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

In C-style order:
0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 

In F-style order:
0, 20, 40, 5, 25, 45, 10, 30, 50, 15, 35, 55,

Modifying Array Elements

The nditer object has another optional parameter op_flags. By default, nditer treats the array as read-only. To modify the array elements while iterating, you must specify the readwrite or writeonly mode.

Example

import numpy as np

a = np.arange(0, 60, 5)
a = a.reshape(3, 4)
print('The original array is:')
print(a)
print('\n')
for x in np.nditer(a, op_flags=['readwrite']):
    x[...] = 2 * x
print('The modified array is:')
print(a)

Output result:

The original array is:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

The modified array is:
[[  0  10  20  30]
 [ 40  50  60  70]
 [ 80  90 100 110]]

Using External Loop

The constructor of the nditer class has a flags parameter, which can accept the following values:

In the following example, the iterator traverses each column and combines them into a one-dimensional array.

Example

import numpy as np
a = np.arange(0, 60, 5)
a = a.reshape(3, 4)
print('The original array is:')
print(a)
print('\n')
print('The modified array is:')
for x in np.nditer(a, flags=['external_loop'], order='F'):
    print(x, end=", ")

Output result:

The original array is:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

The modified array is:
[ 0 20 40], [ 5 25 45], [10 30 50], [15 35 55],

Broadcasting Iteration

If two arrays are broadcastable, the nditer combined object can iterate over them simultaneously. Suppose array a has dimensions 3x4 and array b has dimensions 1x4, then using the following iterator (array b is broadcast to the size of a).

Example

import numpy as np

a = np.arange(0, 60, 5)
a = a.reshape(3, 4)
print('The first array is:')
print(a)
print('\n')
print('The second array is:')
b = np.array([1, 2, 3, 4], dtype=int)
print(b)
print('\n')
print('The modified array is:')
for x, y in np.nditer([a, b]):
    print("%d:%d" % (x, y), end=", ")

Output result:

The first array is:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

The second array is:
[1 2 3 4]

The modified array is:
0:1, 5:2, 10:3, 15:4, 20:1, 25:2, 30:3, 35:4, 40:1, 45:2, 50:3, 55:4,