making this nested for loop more pythonic-CodePudding

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

CodePudding user response：

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 appending.

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.

CodePudding user response：

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]