Search code examples
python-2.7list-comprehensionpython-itertools

python: repeating elements in a list based a predicateHow


I'd like to repeat elements of a list based on a predicate. I tried using the module itertools and a list comprehension

abc = [1,2,3,4,5,6,7,8,9]
result = [ repeat(item,2) if item==3 or item==7 else item for item in abc ]

This doesn't fail at runtime, but the resulting object is not 'flattened' If I print it, I see

[1, 2, repeat(3, 2), 4, 5, 6, repeat(7, 2), 8, 9]

Is it doable with a list comprehension?

Thanks


Solution

  • This works:

    from itertools import repeat
    abc = [1,2,3,4,5,6,7,8,9]
    result = [x for y in (repeat(item,2) if item==3 or item==7 else [item] for item in abc) 
              for x in y]
    
    >>> result
    [1, 2, 3, 3, 4, 5, 6, 7, 7, 8, 9]
    

    The trick here is to put item in its own list [item] and than flatten the now consistently nested list.

    to improve readability, put it in two lines:

    nested_items = (repeat(item,2) if item==3 or item==7 else [item] for item in abc)
    result = [item for nested_item in nested_items for item in nested_item]
    

    Since nested_items is an iterator, there is no extra list created her.