NumPy Array Iteration

NumPy is a powerful library for working with numerical data in Python. One of its key features is the ability to perform efficient element-wise operations on arrays, which makes it particularly useful for data manipulation and scientific computing tasks.

One of the common tasks you might want to perform when working with NumPy arrays is iteration, i.e., looping over the elements of an array and performing some operation on each element. In this post, we'll look at some examples of how to iterate over NumPy arrays and discuss the different options you have for doing so.

Iterating over a 1D NumPy Array

Let's start by creating a simple 1D NumPy array:

import numpy as np

arr = np.array([1, 2, 3, 4, 5])

To iterate over this array, we can simply use a for loop:

for element in arr:
    print(element)

This will print each element of the array on a separate line:

1
2
3
4
5

If you need to access the index of each element as well as the element itself, you can use the enumerate function:

for i, element in enumerate(arr):
    print(f'Index: {i}, Element: {element}')

This will print the index and element for each element in the array:

Index: 0, Element: 1
Index: 1, Element: 2
Index: 2, Element: 3
Index: 3, Element: 4
Index: 4, Element: 5

Iterating over a 2D NumPy Array

Now let's consider a more complex example: a 2D NumPy array. We can create one using the np.array function and passing in a list of lists:

arr = np.array([[1, 2, 3], [4, 5, 6]])

This creates an array with shape (2, 3), i.e., two rows and three columns.

To iterate over this array, we can use a nested for loop:

for row in arr:
    for element in row:
        print(element)

This will print each element of the array on a separate line:

1
2
3
4
5
6

Alternatively, we can use the flat attribute of the array to flatten it and then iterate over the flattened version:

for element in arr.flat:
    print(element)

This will give us the same result as the previous example.

If you need to access the indices of the elements as well, you can use the ndenumerate function:

for index, element in np.ndenumerate(arr):
    print(f'Index: {index}, Element: {element}')

This will print the indices and elements of the array in a format like (row, col), e.g.:

Index: (0, 0), Element: 1
Index: (0, 1), Element: 2
Index: (0, 2), Element: 3
Index: (1, 0), Element: 4
..........