Rethinking Loops As List Comprehensions

·

3 min read

I'm used to writing code that uses loops for everything. Python is stretching my thinking on this. I'm not used to having programming languages written on a higher level with certain procedures written for you. It's like learning a new type of math formula. It took me some time to really understand what was going on, but I was able to solve this problem using them.

Problem (from HackerRank):

Given three integers x, y and z representing the dimensions of a cuboid along with an integer n. Print a list of all possible coordinates given by i, j, k on a 3D grid where the sum of x, y, z is not equal to n.

x = int(input())
y = int(input())
z = int(input())
n = int(input())

listOfPerm = [[i, j, k] for i in range(0,x+1) for j in range(0,y+1) for k in range(0,z+1) if i + j + k != n]

print(listOfPerm)

A good video on this that helped me visualize this was List Comprehension || Python Tutorial || Learn Python Programming from Socratica.

Here's a breakdown of how I determined what was needed. A list comprehension is composed of the following:

  • expression (expr)
  • value (val)
  • collection
<expr> for <values> in <collection>

If it was written as a traditional for loop:

for <values> in <collection>:
    <expr>

This would then need to have a series of nested loops and can get a bit messy.

for i in x:
    for j in y:
        for k in z:

We'd go through each loop to determine what numbers are possible for each variable.

Once we get the resulting values for i, j, k we then also need to check if the sum of all three is not equal to n at the end of the loops.

for i in x:
    for j in y:
        for k in z:
            if ((i+j+k) != n ):
                return [i, j, k]

Fortunately, list comprehensions allows us to write this code much simpler:

<expr> for <values> in <collection> if <condition>
listOfPerm = [[i, j, k] for i in range(0,x+1) for j in range(0,y+1) for k in range(0,z+1) if i + j + k != n]

We want the set of numbers [i, j, k].

For the first number i, let's look at all possibilities from 0 to all possible integer values of x. We do the same with j and i.

Once we have all sets of numbers possible, the condition at the end of the list comprehension will determine which sets of numbers is not equal to n.