Python List Comprehensions

2013-08-22

A friend, who is learning python, asked me about how he might generate a cross-tab-like view of some data he was working on. A straight forward problem that presented a nice opportunity to demonstrate a couple of python's language features to reduce the number of lines of code as well as to make it more readable.

Basic Problem

The end result of the code was a matrix indicating which values where present in each list. The matrix would look something like this:

Some example data in python:

master_list = 'A B C D E F G H'.split()
groups = {
    "GROUP A": 'C E H'.split(),
    "GROUP B": 'A G F D'.split(),
    "GROUP C": 'B C D G'.split(),
    "GROUP D": 'D E F H'.split()
}

Original Code Version

Here's the code that was initially written to produce the table:

matrix = {}

for group_name, group_items in groups.items():
    match_sequence = []
    for item in master_list:
        if item in group_items: 
            match_sequence.append('X')
        else:
            match_sequence.append('-')
    matrix[group_name] = match_sequence

pprint(matrix)

The Ternary Conditional Statement

If we use python ternary conditional statement we can remove the if-else clause and reduce those 4 lines to one.

matrix = {}

for group_name, group_items in groups.items():
    match_sequence = []
    for item in master_list:
        match_sequence.append('X' if item in group_items else '-')
    matrix[group_name] = match_sequence

List Comprehensions

Finally, using a list comprehension we can forgo the for loop structure, and assign the new list to our matrix dict all in one line:

matrix = {}

for group_name, group_items in groups.items():
    matrix[group_name] = ['X' if item in group_items else '-' for item in master_list]

Resources

• Example code
• Python's ternary conditional statement documentation.
• Discussion of C's ternary operator in python.
• Python's list comporehension documentation

Tags: python