Python List Comprehensions
Python is undeniably cool. But one of the coolest things about python is the ability to generate lists from complex arguments over a single line of code without having to resort to multi-line for
loops.
This is well worth learning!
The technique is called List Comprehension, and the syntax is:
[res for value in collection]
res
is the end result, that is, the new list. To create this new list, res
will be evaluated for every item in collection
(which is the starting list).
You can filter what gets included in the new list using if
after the for
loop:
[res for value in collection if <test>]
You can have more than one if
clauses:
[res for value in collection if <test1> and <test2>]
You can also loop over more than one collection!
[res for value1 in collection1 and value2 in collection2]
Here is an example with a list of tuples where we want to filter using one column but transfer the other column into the new list. The collection
here is a list of books and the year they were published:
books = [(‘Frankenstein’, 1818), (‘Twenty Thousand Leagues Under The Sea’, 1870), (‘The Time Machine’, 1895), (‘The War of the Worlds’, 1898), (‘Foundation’, 1942), (‘More Than Human’, 1953), (‘Starship Troopers’, 1959), (‘Solaris’, 1961), (‘Dune’, 1965), (‘The Left Hand of Darkness’, 1969), (‘Neuromancer’, 1984), (‘Snow Crash’, 1992), (‘The Three Body Problem’, 2007)]
To find books that were published before 1900, we just have to whip up this one line list comprehension:
>>> [title for (title, year) in books if year < 1900]
['Frankenstein', 'Twenty Thousand Leagues Under The Sea', 'The Time Machine', 'The War of the Worlds']
How about books published between 1900 and 1960?
>>> [title for (title, year) in books if year > 1900 and year < 1960]
['Foundation', 'More Than Human', 'Starship Troopers']
If you want to also have the year published in the new list:
>>> [(title, year) for (title, year) in books if year < 1900]
[('Frankenstein', 1818), ('Twenty Thousand Leagues Under The Sea', 1870), ('The Time Machine', 1895), ('The War of the Worlds', 1898)]
Elegant!