writing a for loop in list comprehension way to shorten the code

Question

def rahu(big_jug):    

    try_a=list(range(0,int(big_jug/3)+1))    

    for a in try_a:        

        b=(big_jug-3*a)/5        

        if b==int(b):

            return('yes')

    return('no')

Now tried to write in List comprehension way and complete For loop within "Return" ...but getting error and tried in 2-3 ways

def rahu(big_jug): 

    try_a=list(range(0,int(big_jug/3)+1))

    b=(big_jug-3*a)/5

    return ('yes' if b==int(b) for a in try_a  else 'no')

  File "<ipython-input-7-f8ff1bb94105>", line 4
    return ('yes' if b==int(b) for a in try_a  else 'no')
                                 ^
SyntaxError: invalid syntax

I guess the problem is defining ''a'' outside for loop , but how to define this within return with other conditions

Answer 1

Here is some code that does what you want:

def rahu_comprehension(big_jug):
    return next(
        ('yes'
         for a in range(0, int(big_jug / 3) + 1)
         for b in [(big_jug - 3 * a) / 5]
         if b == int(b)
         ),
        'no'
    )

Although this is in multiple lines, this is only for readability, and it works as a single line.

There are a couple of notes I have before I explain this code: Firstly, this is not always a good idea - although some code can be made simpler through the use of list (or other) comprehensions, this code is of the level of complexity that I personally would prefer to use for-loops to make it more readable. Secondly, returning strings 'yes' and 'no' from a function is not good practice. Your function will be more reusable in other places if you return True or False and then if you need a string, you can write 'yes' if rahu(x) else 'no'.

Now, into the code explanation.

The core of the code is the generator expression:

('yes'
 for a in range(0, int(big_jug / 3) + 1)
 for b in [(big_jug - 3 * a) / 5]
 if b == int(b)
 )

Note that this uses parentheses rather than square brackets, which makes it a generator expression not a list comprehension. This means that it doesn't calculate the entire list and return it, but rather only calculates each item when it is needed. In this case we only need to find the first item, since your original code stops calculating as soon as the condition in the if-statement is True, so by using a generator expression, we can stop doing extra work after finding that first value. This generator expression is similar to the following code:

results = []
for a in range(0, int(big_jug / 3) + 1):
    for b in [(big_jug - 3 * a) / 5]:
        if b == int(b):
            results.append('yes')

The outer for a in ... loop is the same as your code. The inner for b in [(big_jug - 3 * a) / 5] is equivalent to a variable assignment, since it is looping through a list of length 1:

results = []
for a in range(0, int(big_jug / 3) + 1):
    b = (big_jug - 3 * a) / 5:
    if b == int(b):
        results.append('yes')

The reason we use this loop is that comprehensions in python don't allow variable assignment, and this is a trick to allow such an assignment.

The next(..., 'no') part of the expression is to take the first item from the generator (which will be a 'yes'), or if the generator has no items, use the default value 'no'.