Suraj Kapoor

Director of Product @WayUp, formerly @Lerer Hippeau. Technologist, Optimistic Contrarian.

List, Set & Dictionary Comprehensions in Python

Comprehensions in python are pretty cool. If you don't use them, I suggest you start right now. Most are familiar with list comprehensions, but set and dictionary comprehensions are available too. Here are simple examples of each.

List Comprehension:

As an example, say you have a list of numbers:

l = [1,2,3,4,5,6,7,8,9,10]  

You can use a basic comprehension to add 100 to each item in the list:

[i+100 for i in l]
>>>[101, 102, 103, 104, 105, 106, 107, 108, 109, 110]

The great things about this is you avoid appending to a new list and creating another object. Less clutter, shorter code.

List Comprehension and Conditional Statements:

You can include conditional statements in the comprehension too. Say you want to add 100 to the even numbers and remove the odd numbers from the list:

[i+100 for i in l if i%2 == 0]
>>>[102, 104, 106, 108, 110]

Now, you need to include an else statement. Let's add 100 to the evens and subtract 50 from the odd:

[i+100 if i%2 == 0 else i-50 for i in l]
>>>[-49, 102, -47, 104, -45, 106, -43, 108, -41, 110]

The placement of the conditional statement changes within the comprehension but it's still fairly straightforward.

Set Comprehension:

Set comprehensions use the curly bracket. Here, I'm adding 100 to evens and subtracting 50 from odds (like the above example):

{i+100 if i%2 == 0 else i-50 for i in l}
>>>set([102, 104, 106, 108, 110, -49, -47, -45, -43, -41])

Remember that elements in a set have no special order.

Dictionary Comprehension:

Similarly, say you have a dictionary as per below, and need to add 10 to each value:

d = {"a":1, "b":2, "c":3}  
{k:d[k]+10 for k in d}
>>>{'a': 11, 'c': 13, 'b': 12}

The important thing to note is that you need to include a colon and provide variables for both the key and value, otherwise the interperter will read it as a set comprehension. Below is an example of a transition from a list to a dictionary comprehension:

l = [1,2,3,4,5,6,7,8,9,10]  
print {i:i for i in l}  
>>>{1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10}

Comprehensions reduce the amount of code quite dramatically and it feels intuitive to keep elements in the same data structure, as opposed to always creating a new one.