CODEWITHSUNDEEP
pythonlist-comprehension
Alex
Alex

Reputation: 93

Multiple if/elif statements in a Python list comprehension

So I'm new to Python and one concept that's taken some getting used to is the list comprehension. I've read that they can improve computation speed when used correctly and that they're something to know well if you're to learn Python the right way.

I'm writing a program that implements a particle aggregation algorithm that involves the accumulation of up to 10^6 particles to a growing cluster. I'm going back through my code to optimize performance anywhere possible, and I have the following function: 

 def step_all_walkers(walkers): 
    updated_positions = []
    for walker in walkers: 
        decision = random.randint(1,4)
        if decision == 1: 
            updated_walker = (min(walker[0]+1, N-1), walker[1])
        elif decision == 2: 
            updated_walker = (max(walker[0]-1, 1), walker[1])
        elif decision == 3: 
            updated_walker = (walker[0], min(walker[1] + 1, N-1))
        elif decision == 4: 
            updated_walker = (walker[0], max(walker[1]-1, 1))

        updated_positions.append(updated_walker)

    return updated_positions

This function makes every particle (or walker, as I call them in the code) take a step of unit length in a random direction, and prevents the particles from walking off an N x N grid. I notice that I am creating and returning a new list updated_positions, and since this list and the input walker list are potentially very large, what I sort of know about list comprehensions tells me this might be a good time to use one. However, in some other posts on this question where there is only one if/else to be evaluated, people responded by saying just use a good ole fashion for loop.

I have a few questions then:

1) Can multiple if/elif statements be done in a list comprehension?

2) Does it make sense to write this for loop as a list comprehension? Are there any advantages to doing so?

My main purpose for asking this question is to build up more intuition for when a list comprehension is appropriate, and also to see if this function can be made more efficient with one.

Upvotes: 0

Views: 380

Answers (2)

DarrylG
DarrylG

Reputation: 17156

A case where list comprehensions should not be used is when the logic in too complex.

For loops are simpler then since they allow:

  • Including comments to explain the code.
  • Use of control flow keywords like continue.
  • Debug of sections of the more easily using logging statements or asserts.
  • Allow others to easily see the complexity of the logic by scanning down the lines

However, as Sayse suggests, we can often simplify the logic with additional data structures and helper functions to allow list comprehension.

List comprehension with no if/else statements (need for logic bypassed)

def step_all_walkers(walkers, N):
  def decision(walker):
    " Helper function for making a decisions "
    # Place decision choices into a data structure
    options = [
      lambda walker: (min(walker[0]+1, N-1), walker[1]),
      lambda walker: (max(walker[0]-1, 1), walker[1]),
      lambda walker: (walker[0], min(walker[1] + 1, N-1)),
      lambda walker: (walker[0], max(walker[1]-1, 1))]

    while True:
      n = random.randint(0, 3)  # use rather than (1, 4) to 
                                # provide proper index into options
      yield options[n](walker)

# Now, list comprehension is straight forward to follow
return [decision(walker) for walker in walkers]

Upvotes: 0

Sayse
Sayse

Reputation: 43300

I would turn that into a lookup dictionary to start with and then you might be able to consider a list comprehension

decisions = {1: lambda walker: (min(walker[0]+1, N-1), walker[1])}
return [decisions[random.randint(1,4)](walker) for walker in walkers]

Upvotes: 1

Related Questions