Using list comprehension and extend

Viewed 78

I have this

rows = self.rows()
aaa = []
for r in range(0, 9, 3):
    bbb = []
    for c in range(0, 9, 3):
        ccc = []
        for s in range(3):
            ccc.extend(rows[r+s][c:c+3])
        bbb.append(ccc)
    aaa.append(bbb)

and it returns this

[
    [
        [5, 0, 0, 0, 6, 0, 0, 2, 9],
        [3, 8, 0, 4, 9, 2, 0, 0, 6],
        [0, 6, 2, 1, 0, 0, 3, 0, 4]
    ],
    [
        [0, 7, 6, 0, 0, 8, 0, 4, 0],
        [0, 4, 0, 2, 0, 5, 0, 3, 1],
        [0, 3, 1, 0, 4, 9, 0, 0, 0]
    ],
    [
        [4, 0, 0, 6, 0, 3, 0, 0, 1],
        [0, 0, 0, 7, 0, 0, 0, 0, 0],
        [0, 5, 3, 2, 0, 8, 0, 0, 6]
    ]
]

which is correct. rows is just a list containing 9 other nested lists, each with exactly 9 integers ranging 0 through 9.

When I try to use list comprehension

[[rows[r+s][c:c+3] for s in range(3) for c in range(0, 9, 3)] for r in range(0, 9, 3)]

I get this

[
    [
        [5, 0, 0],
        [3, 8, 0],
        [0, 6, 2],
        [0, 6, 0],
        [4, 9, 2],
        [1, 0, 0],
        [0, 2, 9],
        [0, 0, 6],
        [3, 0, 4]
    ],
    [
        [0, 7, 6],
        [0, 4, 0],
        [0, 3, 1],
        [0, 0, 8],
        [2, 0, 5],
        [0, 4, 9],
        [0, 4, 0],
        [0, 3, 1],
        [0, 0, 0]
    ],
    [
        [4, 0, 0],
        [0, 0, 0],
        [0, 5, 3],
        [6, 0, 3],
        [7, 0, 0],
        [2, 0, 8],
        [0, 0, 1],
        [0, 0, 0],
        [0, 0, 6]
    ]
]

Clearly I'm doing something wrong, but I can't see what? I've checked other SO questions and they allude to structuring the LC a certain way to stop the innermost list splitting into 9 separate lists, but so far, it's not happening.

3 Answers

Try:

import itertools

[[list(itertools.chain(*[rows[r+s][c:c+3] for s in range(3)])) for c in range(0, 9, 3)] for r in range(0, 9, 3)]

This will give you what you want. This version uses bare Python with no libraries, but seems to be the most concise (of the solutions that work):

[[[rows[r+s][c+i] for s in range(3) for i in range(3)] for c in range(0, 9, 3)] for r in range(0, 9, 3)]

Formatted more readably, this is just:

[[[rows[r+s][c+i] for s in range(3)
                  for i in range(3)]
  for c in range(0, 9, 3)]
 for r in range(0, 9, 3)]

This replaces the slice in the original loop nest with an inner i loop, eliminating the need to concatenate intermediate slices (or to generate them).

You can do something like this,

print([[rows[r+s][c:c+3] * 3 for c in range(0, 9, 3)] for r in range(0, 9, 3)])
Related