Search code examples
pythonconditional-statementslist-comprehension

List comprehension with conditional expression omitting some cases

I have list of bonds between points (as pairs of indexes) and index of a pivot point. I want to list of points bonded to that pivot point irrespective if it is on the first or the second position (I always want the index of the second point to which pivot is bonded in pair).

bonds = [(1,2),(3,4),(5,6),(3,1)]
ipiv  = 1 

bonded_to_pivot = 
[ b[1] for b in bonds if(b[0]==ipiv) ] + 
[ b[0] for b in bonds if(b[1]==ipiv) ]

Can this be done using just one list comprehension in elegant way?

I was looking into these other question about comprehension with conditional expression but I miss something (e.g. else pass) to make it work

Solution

  • Assuming you want the coordinate which doesn't match ipiv, try:

    >>> [b[1] if b[0]==ipiv else b[0] for b in bonds if ipiv in b]
[2, 3]