Show me some cool python list comprehensions

Question

One of the major strengths of python and a few other (functional) programming languages are the list comprehension. They allow programmers to write complex expressions in 1 line. They may be confusing at first but if one gets used to the syntax, it is much better than nested complicated for loops.

With that said, please share with me some of the coolest uses of list comprehensions. (By cool, I just mean useful) It could be for some programming contest, or a production system.

For example: To do the transpose of a matrix mat

>>> mat = [
...        [1, 2, 3],
...        [4, 5, 6],
...        [7, 8, 9],
...       ]

>>> [[row[i] for row in mat] for i in [0, 1, 2]]
[[1, 4, 7], [2, 5, 8], [3, 6, 9]]

Please include a description of the expression and where it was used (if possible).

Answer 1

To do the transpose of a matrix mat:

>>> [list(row) for row in zip(*mat)]
[[1, 4, 7], [2, 5, 8], [3, 6, 9]]

Answer 2

To flatten a list of lists:

>>> matrix = [[1,2,3], [4,5,6]]
>>> [x for row in matrix for x in row]
[1, 2, 3, 4, 5, 6]

Answer 3

A lot of people don't know that Python allows you to filter the results of a list comprehension using if:

>>> [i for i in range(10) if i % 2 == 0]
[0, 2, 4, 6, 8]

Answer 4

I use this all the time when loading tab-separated files with optional comment lines starting with a hash mark:

data = [line.strip().split("\t") for line in open("my_file.tab") \
        if not line.startswith('#')]

Of course it works for any other comment and separator character as well.

Answer 5

As long as you are after functional programming inspired parts of Python, consider map, filter, reduce, and zip----all offered in python.

Answer 6

I currently have several scripts that need to group a set of points into "levels" by height. The assumption is that the z-values of the points will cluster loosely around certain values corresponding to the levels, with large-ish gaps in between the clusters.

So I have the following function:

def level_boundaries(zvalues, threshold=10.0):
    '''Finds all elements z of zvalues such that no other element
    w of zvalues satisfies z <= w < z+threshold.'''
    zvals = zvalues[:]
    zvals.sort()
    return [zvals[i] for i, (a, b) in enumerate(pairs(zvals)) if b-a >= threshold]

"pairs" is taken straight from the itertools module documentation, but for reference:

def pairs(iterable):
    'iterable -> (iterable[n], iterable[n+1]) for n=0, 1, 2, ...'
    from itertools import izip, tee
    first, second = tee(iterable)
    second.next()
    return izip(first, second)

A contrived usage example (my actual data sets are quite a bit too large to use as examples):

>>> import random
>>> z_vals = [100 + random.uniform(-1.5,1.5) for n in range(10)]
>>> z_vals += [120 + random.uniform(-1.5,1.5) for n in range(10)]
>>> z_vals += [140 + random.uniform(-1.5,1.5) for n in range(10)]
>>> random.shuffle(z_vals)
>>> z_vals
[141.33225473458657, 121.1713952666894, 119.40476193163271, 121.09926601186737, 119.63057973814858, 100.09095882968982, 99.226542624083109, 98.845285642062763, 120.90864911044898, 118.65196386994897, 98.902094334035326, 121.2741094217216, 101.18463497862281, 138.93502941970601, 120.71184773326806, 139.15404600347946, 139.56377827641663, 119.28279815624718, 99.338144106822554, 139.05438770927282, 138.95405784704622, 119.54614935118973, 139.9354467277665, 139.47260445000273, 100.02478729763811, 101.34605205591622, 138.97315450408186, 99.186025111246295, 140.53885845445572, 99.893009827114568]
>>> level_boundaries(z_vals)
[101.34605205591622, 121.2741094217216]

Answer 7

I often use comprehensions to construct dicts:

my_dict = dict((k, some_func(k)) for k in input_list)

Note Python 3 has dict comprehensions, so this becomes:

my_dict = {k:some_func(k) for k in input_list}

For constructing CSV-like data from a list of tuples:

data = "\n".join(",".join(x) for x in input)

Not actually a list comprehension, but still useful: Produce a list of ranges from a list of 'cut points':

ranges = zip(cuts, cuts[1:])

Answer 8

If "cool" means crazy, I like this one:

def cointoss(n,t):
    return (lambda a:"\n".join(str(i)+":\t"+"*"*a.count(i) for i in range(min(a),max(a)+1)))([sum(randint(0,1) for _ in range(n)) for __ in range(t)])

>>> print cointoss(20,100)
3:    **
4:    ***
5:    **
6:    *****
7:    *******
8:    *********
9:    *********
10:   ********************
11:   *****************
12:   *********
13:   *****
14:   *********
15:   *
16:   **

n and t control the number of coin tosses per test and the number of times the test is run and the distribution is plotted.