Accelerate Your Python Journey: Mastering Comprehension Filtering

We can also apply filtering to a comprehension. Let’s first do it with filter(), and find all Pythagorean triples whose short sides are numbers smaller than 10. We obviously don’t want to test a combination twice, and therefore we’ll use a trick similar to the one we saw in the previous example:

# pythagorean.triple.p
from math import sqrt
# this will generate all possible pairs
mx = 10
triples = [(a, b, sqrt(a**2 + b**2))
  for a in range(1, mx) for b in range(a, mx)]
# this will filter out all non-Pythagorean triples
triples = list(
  filter(lambda triple: triple[2].is_integer(), triples))
print(triples) # prints: [(3, 4, 5.0), (6, 8, 10.0)]y        

A Pythagorean triple is a triple (a, b, c) of integer numbers satisfying the equation a2 + b2 = c2 .

In the preceding code, we generated a list of three-tuples, triples. Each tuple contains?two?integer numbers (the legs), and the hypotenuse of the Pythagorean triangle whose legs are the first two numbers in the tuple. For example, when a is 3 and b is 4, the tuple will be (3, 4, 5.0), and when a is 5 and b is 7, the tuple will be (5, 7, 8.602325267042627).

After generating all the triples, we need to filter out all those where the hypotenuse is not an integer number. In order to do this, we filter based on float_number. is_integer() being True. This means that, of the two example tuples we just showed you, the one with hypotenuse 5.0 will be retained, while the one with the 8.602325267042627 hypotenuse will be discarded.

This is good, but we don’t like the fact that the triple has two integer numbers and a float — they are all supposed to be integers. Let’s use map() to fix this

# pythagorean.triple.int.p
from math import sqrt
mx = 10
triples = [(a, b, sqrt(a**2 + b**2))
  for a in range(1, mx) for b in range(a, mx)]
triples = filter(lambda triple: triple[2].is_integer(), triples)
# this will make the third number in the tuples integer
triples = list(
  map(lambda triple: triple[:2] + (int(triple[2]), ), triples))
print(triples) # prints: [(3, 4, 5), (6, 8, 10)y        

Notice the step we added. We take each element in triples and we slice it, taking only the first two elements in it. Then, we concatenate the slice with a one-tuple, in which we put the integer version of that float number that we didn’t like. Seems like a lot of work, right? Indeed it is. Let’s see how to do all this with a list comprehension:

# pythagorean.triple.comprehension.p
from math import sqrt
# this step is the same as before
mx = 10
triples = [(a, b, sqrt(a**2 + b**2))
  for a in range(1, mx) for b in range(a, mx)]
# here we combine filter and map in one CLEAN list comprehension
triples = [(a, b, int(c)) for a, b, c in triples if c.is_integer()]
print(triples) # prints: [(3, 4, 5), (6, 8, 10)]y        

That’s much better! It’s clean, readable, and shorter. It’s not quite as elegant as it could have been, though. We’re still wasting memory by constructing a list with a lot of triples that we end up discarding. We can fix that by combining the two comprehensions into one:

# pythagorean.triple.walrus.p
from math import sqrt
# this step is the same as before
mx = 10
# We can combine generating and filtering in one comprehension
triples = [(a, b, int(c))
 for a in range(1, mx) for b in range(a, mx)
 if (c := sqrt(a**2 + b**2)).is_integer()]
print(triples) # prints: [(3, 4, 5), (6, 8, 10)]y        

Now that really is elegant. By generating the triples and filtering them in the same list comprehension, we avoid keeping any triple that doesn’t pass the test in memory. Notice that we used an assignment expression to avoid needing to compute the value of sqrt(a**2 + b**2) twice

We’re going quite fast here, as anticipated in the Summary of , Functions, the Building Blocks of Code. Are you playing with this code? If not, we suggest you do. It’s very important that you play around, break things, change things, and see what happens. Make sure you have a clear understanding of what is going on.

That's it for Today. See you tomorrow.