Iterating Arrays

A for loop over an array walks its first dimension. In 1-D that is each value. In 2-D it is each row, and in 3-D it is each 2-D block. nditer and ndenumerate reach every single value with one loop.

1-D and 2-D loops

Looping a 2-D array gives rows. Add an inner loop to reach each value.

import numpy as np

for p in np.array([1, 2, 3]):
    print(p)

m = np.array([[1, 2, 3], [4, 5, 6]])
for row in m:
    print(row)

for row in m:
    for p in row:
        print(p, end=" ")
print()

Output:

1
2
3
[1 2 3]
[4 5 6]
1 2 3 4 5 6

3-D loops

Looping a 3-D array gives 2-D blocks. Three nested loops reach every value, which gets hard to read quickly.

import numpy as np

t = np.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])
for block in t:
    print(block)

for block in t:
    for row in block:
        for p in row:
            print(p, end=" ")
print()

Output:

[[1 2 3]
 [4 5 6]]
[[ 7  8  9]
 [10 11 12]]
1 2 3 4 5 6 7 8 9 10 11 12

nditer

np.nditer(arr) gives every value with one loop, whatever the number of dimensions. For a normal array it reads row by row.

import numpy as np

t = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
for p in np.nditer(t):
    print(p, end=" ")
print()

Output:

1 2 3 4 5 6 7 8

ndenumerate

np.ndenumerate(arr) gives the index tuple and the value together. Here m is the index and p is the value.

import numpy as np

grid = np.array([[7, 3], [5, 9]])
for m, p in np.ndenumerate(grid):
    print(m, p)

Output:

(0, 0) 7
(0, 1) 3
(1, 0) 5
(1, 1) 9

When you do not need a loop

Most array work is a vector operation. A loop is fine for printing or for logic that depends on each value. For totals and element math, the array method is shorter and much faster.

import numpy as np

arr = np.arange(1, 101)
total = 0
for p in arr:
    total += p
print(total)
print(arr.sum())

Output:

5050
5050

Key takeaways

• A loop walks the first dimension: values, rows, or blocks

• nditer reaches every value with one loop

• ndenumerate gives each index tuple with its value

• Prefer array methods such as sum over a manual loop

• Next: Arithmetic and Broadcasting

• Previous: Splitting Arrays