Search code examples
pythonlistoperation

How do I perform multiple operations in a list comprehension


L = [random.randint(0,50) for i in range(5) random.randint(0,12) for i in range(2)]

How do I get it to pick 5 random numbers between (0,50), then 2 random numbers between(0,12)?


Solution

  • You can vary the second argument to randint() based on the value of i:

    [randint(0, 50 if i < 5 else 12) for i in range(7)]
    

    The 50 if i < 5 else 12 expression will change what is passed to random.randint() for the last two iterations.

    There are many more variations you can spell this in. List comprehensions are a bunch of loops and if filters that repeatedly execute the expression at the front. There are lots of ways to spell vary the arguments to a function call based on the iteration values in expressions.

    For example, you could record those arguments in functools.partial() objects:

    from functools import partial
    from random import randint
    
    rint50 = partial(randint, 0, 50)
    rint12 = partial(randint, 0, 12)
    [rint() for rint in [rint50] * 5 + [rint12] * 2]
    

    The possibilities are endless. Lambdas, randint(0, upperbound), randint(*args), a function that'll vary its results depending on how often it has been called, etc. But I wouldn't argue that any of these are actually more readable or understandable.

    For this case, with just 7 values, I'd just concatenate the two lists:

    [randint(0, 50) for _ in range(5)] + [randint(0, 12) for _ in range(2)]
    

    as it's just cleaner and more readable. The small performance cost of creating a 3rd list that contains the results of the two list comprehensions is negligible here.