making this nested for loop more pythonic

Question

I have a clunky (but working) piece of code as shown:

plus_cords = []
for i in range(len(pluses)):
    plus_cords.append([ [pluses[i][0], pluses[i][1]] ])
    for j in range(1, pluses[i][2] + 1):
        plus_cords[i].append([pluses[i][0] - j, pluses[i][1]])
        plus_cords[i].append([pluses[i][0] + j, pluses[i][1]])
        plus_cords[i].append([pluses[i][0], pluses[i][1] - j])
        plus_cords[i].append([pluses[i][0], pluses[i][1] + j])

where 'pluses' is a list of a list of 3 integers.

pluses = [[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 3, 0], [0, 4, 0], [0, 5, 0], [1, 0, 0], [1, 4, 0], [2, 0, 0], [2, 1, 0], [2, 2, 0], [2, 3, 0], [2, 4, 0], [2, 4, 1], [2, 5, 0], [3, 0, 0], [3, 1, 0], [3, 4, 0], [4, 0, 0], [4, 1, 0], [4, 2, 0], [4, 3, 0], [4, 4, 0], [4, 5, 0]]

I'm looking for ideas on how this can be made more readable and efficient, basically more "pythonic".

Thank you in advance

Answer 1

In general it's preferable to iterate over the elements of a list rather than iterating over it by index, and to build a new list via comprehension rather than iteratively append ing.

plus_cords = [[
    [p0, p1], [p0 - j, p1], [p0 + j, p1], [p0, p1 - j], [p0, p1 + j]
] for [p0, p1, p2] in pluses for j in range(1, p2 + 1)]

Not positive I got it right since I don't have sample data for pluses to test with, but that should be in the general ballpark.

Answer 2

Similar to the answer by Samwise, but broken down into separate steps, and actually producing the same result as your original code:

First, of course, we can iterate the elements of pluses directly, instead of using an index, and unpack them to their three constituting values, already making the code a good deal more readable:

plus_cords = []
for p0, p1, p2 in pluses:
    cords = [ [p0, p1] ]
    for j in range(1, p2 + 1):
        cords.append([p0 - j, p1])
        cords.append([p0 + j, p1])
        cords.append([p0, p1 - j])
        cords.append([p0, p1 + j])
    plus_cords.append(cords)

Then, we can try to turn the inner loop into a list comprehension. This is a bit tricky, as we have to use a nested generator and unpack that into the list comprehension to differentiate between the single first element and the rest:

plus_cords = []
for p0, p1, p2 in pluses:
    cords = [[p0, p1], *(x for j in range(1, p2 + 1)
                           for x in ([p0 - j, p1], [p0 + j, p1], [p0, p1 - j], [p0, p1 + j]))]
    plus_cords.append(cords)

Once we have that, we can make the outer loop a list comprehension, too.

plus_cords = [[[p0, p1], *(x for j in range(1, p2 + 1)
                             for x in ([p0 - j, p1], [p0 + j, p1], [p0, p1 - j], [p0, p1 + j]))]
              for p0, p1, p2 in pluses]

Alternatively, instead of the second step, you could settle for a set of tuples (to prevent duplicates) and get the [p0, p1] case with starting with j=0 .

plus_cords = [{x for j in range(p2 + 1)
                 for x in ((p0 - j, p1), (p0 + j, p1), (p0, p1 - j), (p0, p1 + j))}
              for p0, p1, p2 in pluses]