Iterating Over Arrays

Iterating

Iterating over arrays is sometimes necessary, but try to avoid it when you can. Think of it like using a calculator vs doing mental math - vectorization is usually better.

Basic Iteration with nditer

nditer is the standard way to iterate over arrays.

import numpy as np

arr = np.array([[1, 2, 3], [4, 5, 6]])
for x in np.nditer(arr):
    print(x, end=' ')
print()
    

nditer flattens the array and iterates through each element. Simple and clean.

Enumerate with ndenumerate

ndenumerate gives you both index and value.

import numpy as np

arr = np.array([[1, 2, 3], [4, 5, 6]])
for idx, val in np.ndenumerate(arr):
    print(f"Index {idx}: {val}")
    

Here is the thing - idx is a tuple for multi-dimensional arrays. (0, 0) means row 0, column 0.

When to Avoid Loops

Vectorized operations are almost always faster than loops.

import numpy as np
import time

arr = np.arange(1000000)

start = time.time()
result1 = arr ** 2
print(f"Vectorized: {time.time() - start:.4f}s")

start = time.time()
result2 = np.array([x**2 for x in arr])
print(f"List comprehension: {time.time() - start:.4f}s")
    

One thing that confused me at first was when loops are okay. Use them for small arrays or when vectorization is impossible.

Key Takeaways

• nditer iterates through array elements
• ndenumerate gives index and value pairs
• Prefer vectorized operations over loops
• Loops are okay for small arrays or complex logic