Search code examples
pythontupleslist-comprehension

Unpack tuples in list comprehension with a condition


I'd like to build one of these lists of tuples:

  • (a, 0), (-a, 0) (b, 0), (-b, 0)
  • (0, a), (0, -a) (0, b), (0, -b)

from scalars a and b.

based on a condition:

  • c = a > b

This is my attempt:

a = 5
b = 2
c = a > b

# Try build two tuples per element, e.g. (5, 0), (-5, 0) (2, 0), (-2, 0)

# This syntax is illegal
#f2 = [(m,0), (-m,0) if c else (0,m), (-0,-m) for m in (a,b)]

# This syntax works but creates tuples of tuples
f2 = [tuple(((m,0), (-m,0))) if c else tuple(((0,m), (-0,-m))) for m in (a,b)]
print(*f2) # ((5, 0), (-5, 0)) ((2, 0), (-2, 0))

# This syntax is illegal
#f3 = [*tuple(((m,0), (-m,0))) if c else *tuple(((0,m), (-0,-m))) for m in (a,b)]
#print(*f3)

f2 builds a list of two tuples of two tuples: ((5, 0), (-5, 0)) ((2, 0), (-2, 0)).
Using * operator in f3 to unpack the outer tuples triggers a syntax error.

What is the correct syntax?


Also I don't understand why f2 is ((5, 0), (-5, 0)) ((2, 0), (-2, 0)), where the outer tuples are not separated by a ,?


Solution

  • You can do it with lambda functions.

    x1 = lambda x : (x,0)
    y1 = lambda x : (-x,0)
    x2 = lambda x : (0,x)
    y2 = lambda x : (-0,-x)
    f2 = [f1(m) if c else f2(m) for m in (a,b) for f1,f2 in zip((x1,y1),(x2,y2))] 
    

    Output:
    [(5, 0), (-5, 0), (2, 0), (-2, 0)]
    A little bit over-engineered but makes sense.

    EDIT
    Also you can use chain from itertools to put 2 items in a single list comprehension.

    from itertools import chain
    f2 = list(
        chain.from_iterable(
            ((m,0), (-m,0)) if c else ((0,m), (-0,-m)) for m in (a,b)
        ))
    

    SEE