Learning Objectives – after this chapter you will be able to:
- Loop over 1-D, 2-D and 3-D arrays with
for. - Visit every element of any array with
np.nditer(). - Get indexes and values together with
np.ndenumerate(). - Explain why vectorised operations are preferred over loops.
Iterating with a for Loop
Iterating means going through elements one by one. A for loop over an array moves along the first axis: for a 1-D array you get elements, for a 2-D array you get rows, for a 3-D array you get 2-D blocks.
import numpy as np
a1 = np.array([1, 2, 3])
for x in a1:
print(x, end=" ")
print()
a2 = np.array([[1, 2, 3], [4, 5, 6]])
for row in a2:
print("row:", row)
for row in a2: # nested loop for every element
for x in row:
print(x, end=" ")
print()Output
1 2 3 row: [1 2 3] row: [4 5 6] 1 2 3 4 5 6
3-D array needs three nested loops
import numpy as np
a3 = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
for block in a3:
for row in block:
for x in row:
print(x, end=" ")
print()Output
1 2 3 4 5 6 7 8
np.nditer() – One Loop for Any Dimension
np.nditer() visits every element of an array of any dimension with a single loop. It can also step through elements and change their data type while iterating.
import numpy as np
a3 = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
for x in np.nditer(a3):
print(x, end=" ")
print()
a2 = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
for x in np.nditer(a2[:, ::2]): # every 2nd column
print(x, end=" ")
print()
for x in np.nditer(np.array([1, 2, 3]), flags=["buffered"], op_dtypes=["float64"]):
print(x, end=" ")
print()Output
1 2 3 4 5 6 7 8 1 3 5 7 1.0 2.0 3.0
Changing values while iterating
import numpy as np
arr = np.array([[1, 2], [3, 4]])
with np.nditer(arr, op_flags=["readwrite"]) as it:
for x in it:
x[...] = x * 10 # modify in place
print(arr)Output
[[10 20] [30 40]]
np.ndenumerate() – Index and Value
import numpy as np
arr = np.array([[10, 20], [30, 40]])
for index, value in np.ndenumerate(arr):
print(index, int(value))Output
(0, 0) 10 (0, 1) 20 (1, 0) 30 (1, 1) 40
Prefer Vectorised Operations
Python loops over large arrays are slow. Most tasks can be done with vectorised NumPy operations that act on the whole array at once and run in compiled code.
import numpy as np
arr = np.array([3, 8, 1, 9, 4])
# Loop version
total = 0
for x in arr:
total += x * x
print(total)
# Vectorised version - same result, faster on big arrays
print((arr ** 2).sum())
print(arr * 2 + 1)Output
171 171 [ 7 17 3 19 9]
Key Takeaways
- A for loop moves along the first axis; nested loops reach every element.
np.nditer()visits every element of any array in one loop;np.ndenumerate()adds the index.- Use vectorised operations instead of loops whenever possible.
Video : Array Iterating
Question / Answer (15 questions)
Answer : Iterating means visiting the elements of an array one by one, usually with a for loop, to read or process them.
Answer : The loop moves along the first axis, so each iteration gives one row as a 1-D array.
Answer : Three: one for the blocks, one for the rows within each block and one for the elements within each row.
Answer : np.nditer() is an efficient iterator that visits every element of an array of any dimension using a single loop, avoiding nested loops.
Answer : Pass a slice to nditer: for x in np.nditer(arr[:, ::2]): …
Answer : Use the op_dtypes argument together with flags=["buffered"], for example np.nditer(arr, flags=["buffered"], op_dtypes=["float64"]). Buffering provides temporary space for the converted values.
Answer : Open the iterator with op_flags=["readwrite"] inside a with block and assign through x[...] = new_value.
Answer : It yields pairs of (index tuple, value) for every element, e.g. ((0, 1), 20), so you know the position of each element.
Answer : The loop variable is a separate value; reassigning it only changes the local name. To modify the array you must assign through an index (arr[i] = …) or use nditer with readwrite and x[...].
Answer : Vectorisation means performing an operation on whole arrays at once (e.g. arr * 2) instead of looping over elements in Python. NumPy runs these operations in optimised compiled code.
Answer : They are much faster on large arrays, use shorter and more readable code, and reduce chances of indexing mistakes.
Answer : (arr ** 2).sum(). For [3, 8, 1, 9, 4] the result is 171.
Answer : For a normally created array it visits elements row by row (C order): all elements of the first row, then the second row, and so on.
Answer : For small arrays, for logic that cannot be expressed with vectorised operations, or when processing items with side effects such as printing or writing files.
Answer : 1 2 3 4
MCQ Test : Check Your Understanding (5 questions)
Select an option and press Check Answer, or answer all questions and press Submit Test to see your score. Each answer comes with an explanation.
Bibliography/References
The notes, questions and MCQs on this page are original explanations written for this course. The works listed below were consulted and are cited for attribution and further reading. Product names, logos and trademarks belong to their respective owners. Short quotations, where used, are marked and attributed.
- NumPy Developers. (n.d.). Iterating over arrays. NumPy Reference. https://numpy.org/doc/stable/reference/arrays.nditer.html
- NumPy Developers. (n.d.). NumPy quickstart. NumPy User Guide. https://numpy.org/doc/stable/user/quickstart.html
- NumPy Developers. (n.d.). NumPy: the absolute basics for beginners. NumPy User Guide. https://numpy.org/doc/stable/user/absolute_beginners.html
- VanderPlas, J. (2016). Python Data Science Handbook. O’Reilly Media. (Free online edition.) https://jakevdp.github.io/PythonDataScienceHandbook/
- Harris, C. R., Millman, K. J., van der Walt, S. J., et al. (2020). Array programming with NumPy. Nature, 585, 357–362. https://doi.org/10.1038/s41586-020-2649-2