Iterate Pythonically
Problem to Solve
I want to accomplish various iteration-related tasks while keeping my code simple and easy.
When iterating over the elements in a collection, it's easy to fall back on counter-based iteration like one would use in an old-fashioned language such as fortran, C, or java. Don't do that! Python has a different conceptualization of "iteration", and it's more powerful, convenient, and compact. Some situations do require the use of a counter variable and index-based collection access, but they're fairly rare.
It's also worth remembering that many tasks that require iteration in lower-level languages can be done with a single function or method call in python. For example, to find the sum of the elements in a list, you can do sum(some_list)
instead of adding them up one at a time via iteration.
What follows is a few example recipes to illustrate the "pythonic way" to iterate over collections.
A Few Recipes
do something with every element of a tuple, list, or set
for element in some_list_set_or_collection:
# do something with element
iterate over a sequence's elements in reverse order
for element in reversed(some_sequence):
# do something with element
keep a counter variable while iterating
Sometimes you do need to count the iterations as you iterate through a sequences. That does not mean you should reach for for i in range(...)
and access the elements by index! Python provides the enumerate()
function specifically for this:
for i, element in enumerate(some_sequence):
print(f"Element number {i} is {element}.")
iterating over corresponding items of two (or more) sequences
list_of_sums =[]
for element1, element2 in zip(sequence1, sequence2):
list_of_sums.append(element1 + element2)
iterating over corresponding items of two (or more) sequences while keeping a counter
for i, (x, y) in enumerate(zip(x_values, x_values)):
print(f"The coordinates of point {i} are ({x}, {y}).")
iterating over all combinations of elements from two sequences
The old-school way:
# Not very pythonic:
for x in x_values:
for y in y_values:
print(f"The coordinates of point ({x}, {y}).")
The pythonic way:
from itertools import product
for x, y in product(x_values, y_values):
print(f"The coordinates of point ({x}, {y}).")
It's only the same number of lines the first time you use it; after that, the import statement doesn't need to be repeated. And in general, the less nested your code is, the easier it is to comprehend.
iterating over a sequence in multi-element chunks
from itertools import batched
for tuple_of_n_successive_elements in batched(some_sequence, n):
# do something with each successive set of `n` elements
Bonus Skill: Comprehensions
For tasks that require a short, simple iterative loop to produce a new list, set, or dict, python provides a special syntax called a comprehension. Comprehensions are a compact way to express the same thing you would do with a loop, but in a single line of code. They are often more readable than the equivalent loop, and they can be faster too.
calculate values of a new list from corresponding elements of an old list
old_list = [1.2, 3.4, 7.9, 12.4] # just an arbitrary example
squares_of_old = [x**2 for x in old_list]
roots_of_old = [math.sqrt(x) for x in old_list]
roots_of_all = [math.sqrt(n) for n in range(1, 100)]
filter a list to keep only some elements
data = [random.uniform(-1, 1) for _ in range(10)] # just fake data
positive_data = [x for x in data if x > 0]
divisible_by_7 = [x for x in range(1, 100) if x % 7 == 0]
calculate values of a new list from only some values in an old list
data_sqrts_1 = [math.sqrt(x) for x in data if x > 0] # omit negative values
data_sqrts_2 = [math.sqrt(x) if x > 0 else 0 for x in data] # replace negative values with 0
Comprehensions can also make sets, dictionaries, and a special kind of object called a generator, just by changing the syntax slightly.
combine a keywords list and a values list into a dictionary
keys = ['mass', 'charge', 'spin']
values = [9.11e-31, -1.6e-19, 1/2]
dict1 = {key: val for key, val in zip(keys, values)}
Note the curly-braces. List comprehensions use square brackets; set comprehensions use curly braces; and dict comprehensions use curly braces with a colon between the key and value.