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

If the code works and you're looking for advice on improving it, Code Review is the appropriate place. But see codereview.meta.stackexchange.com/questions/5777/… first. — Barmar
– Barmar, Commented Sep 25, 2021 at 19:24
Noted, I'll keep this in mind. I have also posted it over there. Thanks. — Muffi
– Muffi, Commented Sep 25, 2021 at 20:11

tobias_k · Accepted Answer · 2021-09-25 20:34:46Z

1

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]

answered Sep 25, 2021 at 20:34

tobias_k

83.1k12 gold badges130 silver badges186 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Muffi Over a year ago

Thank you Tobias. For the outer for loop in the first step, what would be the difference between "for p0, p1, p2 in pluses" and "for [p0, p1, p2] in pluses". I get the correct result either way. For step 2 onwards, I'll definitely have to do some reading to understand. Would you have any resources for this particular content?

tobias_k Over a year ago

@Muffi (1) There is not really a difference. (2) Best just try those in an interactive python shell: simple list literal [1, 2, 3]; unpacking another iterable into a list literal [1, *(2, 3)]; list comp. [x*2 for x in range(3)]; list comp with 2 loops: [x+y for x in range(3) for y in range(3)], then try to combine those in different ways. What you see in step 2 is a generator expression with 2 for-loops that is unpacked into a list literal.

Samwise · Accepted Answer · 2021-09-25 19:49:21Z

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

edited Sep 25, 2021 at 19:49

answered Sep 25, 2021 at 19:29

Samwise

72.1k3 gold badges36 silver badges52 bronze badges

4 Comments

wjandrea Over a year ago

This has a syntax error. It looks like the for j part is in the wrong place

Muffi Over a year ago

Yes, I did try this on my end, but i'm getting a syntax error. Although this does look much more readable, thanks. I'll try to debug the syntax error and post over here.

Samwise Over a year ago

ah yes, we're doing one list per element of pluses, which means we want for p... in pluses for j in range(...). Updated.

Muffi Over a year ago

That does get rid of the syntax error, but doesn't lead to the expected output. I'll add an example of pluses to the main question for your reference

Collectives™ on Stack Overflow

making this nested for loop more pythonic

2 Answers 2

2 Comments

4 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

4 Comments

Your Answer

Sign up or log in

Post as a guest

Related