Search code examples
pythonlist-comprehension

List comprehension: Returning two (or more) items for each item


Is it possible to return 2 (or more) items for each item in a list comprehension?

What I want (example):

[f(x), g(x) for x in range(n)]

should return [f(0), g(0), f(1), g(1), ..., f(n-1), g(n-1)]

So, something to replace this block of code:

result = list()
for x in range(n):
    result.add(f(x))
    result.add(g(x))

Solution

  • >>> from itertools import chain
    >>> f = lambda x: x + 2
    >>> g = lambda x: x ** 2
    >>> list(chain.from_iterable((f(x), g(x)) for x in range(3)))
    [2, 0, 3, 1, 4, 4]
    

    Timings:

    from timeit import timeit
    
    f = lambda x: x + 2
    g = lambda x: x ** 2
    
    def fg(x):
        yield f(x)
        yield g(x)
    
    print timeit(stmt='list(chain.from_iterable((f(x), g(x)) for x in range(3)))',
                 setup='gc.enable(); from itertools import chain; f = lambda x: x + 2; g = lambda x: x ** 2')
    
    print timeit(stmt='list(chain.from_iterable(fg(x) for x in range(3)))',
                 setup='gc.enable(); from itertools import chain; from __main__ import fg; f = lambda x: x + 2; g = lambda x: x ** 2')
    
    print timeit(stmt='[func(x) for x in range(3) for func in (f, g)]',
                 setup='gc.enable(); f = lambda x: x + 2; g = lambda x: x ** 2')
    
    
    print timeit(stmt='list(chain.from_iterable((f(x), g(x)) for x in xrange(10**6)))',
                 setup='gc.enable(); from itertools import chain; f = lambda x: x + 2; g = lambda x: x ** 2',
                 number=20)
    
    print timeit(stmt='list(chain.from_iterable(fg(x) for x in xrange(10**6)))',
                 setup='gc.enable(); from itertools import chain; from __main__ import fg; f = lambda x: x + 2; g = lambda x: x ** 2',
                 number=20)
    
    print timeit(stmt='[func(x) for x in xrange(10**6) for func in (f, g)]',
                 setup='gc.enable(); f = lambda x: x + 2; g = lambda x: x ** 2',
                 number=20)
    

    2.69210777094

    3.13900787874

    1.62461071932

    25.5944058287

    29.2623711793

    25.7211849286