How to use set comprehensions in Python

What are set comprehensions? How do they differ from list comprehensions, and when are they useful.

How to use set comprehensions in Python

We've talked about list comprehension before.

Comprehensions are great, and we should use them more. But comprehensions are not limited to lists. We can use comprehensions to build sets as well.

What is a set?Permalink

set is an unordered collection of unique items. It can't contain the same items.

It is useful when we want to make sure that we do not store duplicate items -- and we don't care about order.

We define a set the same way as a list, except we use curly brackets {} instead of square brackets []:

unique_numbers = {1, 1, 2, 2, 3, 3}
unique_numbers
# outputs {1, 2, 3}

Above we wrote every number twice when defining our set, but the resulting set only contains each number once. That is what set does for us.

Set comprehensionsPermalink

What a set comprehension does? It creates a set using existing iterable as an input.

Set comprehensions work the same as list comprehensions. The only difference is that they use curly brackets {} instead of square brackets []:

Here's a minimal set comprehension that uses some_iterable as an input and creates a new set:

{thing for thing in some_iterable}

Let's see it in action with some data.

Say we have a list of numbers where some numbers are there twice. We want to square all numbers in this list - but only keep unique results.

We could use list comprehension to build a list of squares:

numbers = [2, 4, 3, 4, 5, 6, 5] # 4 & 5 are in this list twice

squares = [n * n for n in numbers]
squares
# outputs: [4, 16, 9, 16, 25, 36, 25]
# squares 16 and 25 are there twice!

But our list of results contains numbers 16 & 26 twice. Not what we wanted. To fix this, we can create a set from the list:

unique_squares = set(squares)
unique_squares
# outputs: {4, 36, 9, 16, 25}

But this is not ideal because:

  • We first create the list and then create the set.
  • So, we do two operations.
  • And allocate memory twice for all items.

Instead, we could create a set straight away using set comprehension:

unique_squares = {n * n for n in numbers}
unique_squares
# outputs: {4, 36, 9, 16, 25}

And that's all there is to it. We just used different brackets. Our code is now shorter, faster and more memory efficient.

As with list comprehension, we could rewrite a set comprehension as a for loop :

unique_squares = set()

for n in numbers:
    unique_squares.add(n*n)

Here we can see how the above for loop translates to a set comprehension: How set comprehension translates to the for loop

And of course, as with list comprehension, we can also use if to filter what we want to process:

{n * n for n in numbers if n > 3}
# outputs: {36, 16, 25}

SummaryPermalink

Set comprehensions work in the same way as a list comprehensions but:

  • they create a set (not a list)
  • they start and end with curly brackets {} (instead of square brackets [])

Happy coding!

You might also like

Better Python apps in AWS with stelvio.dev

Deep dives and quick insights into cloud architecture.

    We won't send you spam. Unsubscribe at any time.