Un peu de math
An overview of the awesome `zip` generator in Python.
Iterating over many lists at the same time in Python with `zip`
2018-11-23
When programming it is often necessary to iterate over two (or more) lists. This post will describe zip: a fantastic tool for doing this in a clear and concise way.
As a running example, let us assume we would like to check if 3 numbers $a, b, c$ form a pythagorean triplet, in other words:
$$ a ^ 2 + b ^ 2 = c ^ 2 $$For the purpose of this example let us start by randomly creating 3 lists of numbers:
import random
random.seed(0)
upper_value = 50
number_of_values = 10 ** 5
list_of_a_values = [random.randint(1, upper_value + 1) for _ in range(number_of_values)]
list_of_b_values = [random.randint(1, upper_value + 1) for _ in range(number_of_values)]
list_of_c_values = [random.randint(1, upper_value + 1) for _ in range(number_of_values)]
Here are our first 5 values of $a$:
list_of_a_values[:5]
A common approach¶
One commone approach would be to create a dummy variable that will keep track of a position that goes through all 3 lists:
pythagorian_triplets = []
for i in range(number_of_values):
if list_of_a_values[i] ** 2 + list_of_b_values[i] ** 2 == list_of_c_values[i] ** 2:
pythagorian_triplets.append(
(
list_of_a_values[i],
list_of_b_values[i],
list_of_c_values[i],
)
)
pythagorian_triplets
Zipping our lists together¶
Another approach is to use python's zip command, this will essentially "zip" the three lists together
pythagorian_triplets = []
for item in zip(list_of_a_values, list_of_b_values, list_of_c_values):
if item[0] ** 2 + item[1] ** 2 == item[2] ** 2:
pythagorian_triplets.append(item)
pythagorian_triplets
What zip is doing here is taking a list of inputs, in this case list_of_a_values, list_of_b_values, list_of_c_values (it can take any number of inputs) and returning a python object that can be iterated over. Each iteration takes an element from each of the inputs.
We can use python unpacking and list comprehensions to write this in a more compact way:
pythagorian_triplets = [(a, b, c) for a, b, c in zip(list_of_a_values, list_of_b_values, list_of_c_values)
if a ** 2 + b ** 2 == c ** 2]
pythagorian_triplets
Unpacking and zipping¶
Using numpy we can in fact make this slightly more efficient as we can generate a large 3 by N array of random integers (note that numpy's random number generator does not follow the same seeded path as python's standard lib):
import numpy as np
np.random.seed(0)
random_integers = np.random.randint(1, upper_value, (3, number_of_values))
random_integers
We can now use python unpacking with the * command to pass each row of our array directly to zip without having to name them:
pythagorian_triplets = [(a, b, c) for a, b, c in zip(*random_integers)
if a ** 2 + b ** 2 == c ** 2]
pythagorian_triplets
zip is one of those super handy python commands that once you've gotten the hang off you never want to be without, enumerate is another such tool.
A blog about programming (usually scientific python), mathematics (usually game theory) and learning (usually student centred pedagogic approaches).
Source code: drvinceknight Twitter: @drvinceknight Email: [email protected] Powered by: Python mathjax highlight.js Github pages Bootstrap