Is it possible to use "zip()" with a step parameter?

Question

I would like to extract a very specific portion of a 2D array in Python using the zip() method (and avoid messy for loop logic). I'd like to use zip to achieve something like this:

>>> sub_matrix = list(zip([*grid[0:3]]*3))

# Desired output example (Option 1)
[".","4",".", ".",".","4",".",".","."]

# Desired output example (Option 2)
[[".","4","."],  
[".",".","4"],
[".",".","."]]

I am working with the 2D array below in Python on an interview practice problem.

grid = [[".","4",".",".",".",".",".",".","."],
    [".",".","4",".",".",".",".",".","."],
    [".",".",".","1",".",".","7",".","."],
    [".",".",".",".",".",".",".",".","."],
    [".",".",".","3",".",".",".","6","."],
    [".",".",".",".",".","6",".","9","."],
    [".",".",".",".","1",".",".",".","."],
    [".",".",".",".",".",".","2",".","."],
    [".",".",".","8",".",".",".",".","."]]

Part of solving the problem involves ensuring each 3 x 3 "region" in a sudoku game contains legal values. I'd like to use zip() to quickly extract a 3 x 3 portion of the matrix. For example, the top-left region would cause the tests to fail because it contains 4 twice.

I know I can subset the grid to get the first three rows as follows:

    >>> sub_grid = grid[0:3]
    >>> print(sub_grid)
    [['.', '4', '.', '.', '.', '.', '.', '.', '.'], 
    ['.', '.', '4', '.', '.', '.', '.', '.', '.'], 
    ['.', '.', '.', '1', '.', '.', '7', '.', '.']]

I modified the printing a little bit to make it obvious, but at this point, I'd like to zip the three arrays using a 'step' of 3, so that each new array will zip 3 values from each array before moving on to the next one.

In the Python3 docs on zip there is an excerpt on how I think this can be done, but I cannot get the desired output.

The left-to-right evaluation order of the iterables is guaranteed. This makes possible an idiom for clustering a data series into n-length groups using zip(*[iter(s)]*n).

(For posterity, question is from CodeFights will be hidden until unlocked)

Any help is greatly appreciated. Thanks.

In numpy, this would be `grid[0:3, 0:3]`. – cs95 Jun 26 '17 at 16:40 — cs95, Jun 26 '17 at 16:40

score 4 · Accepted Answer · answered Jun 25 '17 at 17:18

4

No zip but [row[:3] for row in grid[:3]]

answered Jun 25 '17 at 17:18

woodpav

1,917
2
13
26

Very nice answer, thank you. I'll wait to see if there is a zip related answer before ticking the check-box. Cheers. – Scott Skiles Jun 25 '17 at 17:25

score 2 · Answer 2 · answered Jun 25 '17 at 17:33

zip does not support stepping through nested structures in that manner. Another alternative to the other answer is to use a map and extract those slices.

from operator import itemgetter
list(map(itemgetter(0, 1, 2), grid[0:3]))

or, if you prefer:

list(map(lambda x: x[0:3], grid[0:3]))

If you're using python2, you can drop the extra list(...).

As for what technique to use and why, take a look at this thread on SO.

Output:

[('.', '4', '.'), ('.', '.', '4'), ('.', '.', '.')]

fafl · Answer 3 · 2017-07-17T13:07:47.630

TL;DR

To get the blocks run this one-liner:

[reduce(lambda a, b: a+b, item) for l in [zip(*row) for row in zip(*[iter([zip(*[iter(row)]*3) for row in grid])]*3)] for item in l]

Step by step explanation

Let's first look at how to use zip:

chunks = [zip(*[iter(row)]*3) for row in grid]

[
    [('.', '4', '.'), ('.', '.', '.'), ('.', '.', '.')],
    [('.', '.', '4'), ('.', '.', '.'), ('.', '.', '.')],
    [('.', '.', '.'), ('1', '.', '.'), ('7', '.', '.')],
    [('.', '.', '.'), ('.', '.', '.'), ('.', '.', '.')],
    [('.', '.', '.'), ('3', '.', '.'), ('.', '6', '.')],
    [('.', '.', '.'), ('.', '.', '6'), ('.', '9', '.')],
    [('.', '.', '.'), ('.', '1', '.'), ('.', '.', '.')],
    [('.', '.', '.'), ('.', '.', '.'), ('2', '.', '.')],
    [('.', '.', '.'), ('8', '.', '.'), ('.', '.', '.')]
]

You see how each row got split into chunks of size three. The rows should also be split, so we run:

blocks = zip(*[iter(chunks)]*3)

[
    (
        [('.', '4', '.'), ('.', '.', '.'), ('.', '.', '.')],
        [('.', '.', '4'), ('.', '.', '.'), ('.', '.', '.')],
        [('.', '.', '.'), ('1', '.', '.'), ('7', '.', '.')]
    ), (
        [('.', '.', '.'), ('.', '.', '.'), ('.', '.', '.')],
        [('.', '.', '.'), ('3', '.', '.'), ('.', '6', '.')],
        [('.', '.', '.'), ('.', '.', '6'), ('.', '9', '.')]
    ), (
        [('.', '.', '.'), ('.', '1', '.'), ('.', '.', '.')],
        [('.', '.', '.'), ('.', '.', '.'), ('2', '.', '.')],
        [('.', '.', '.'), ('8', '.', '.'), ('.', '.', '.')]
    )
]

This gives us the blocks. For extraction it makes sense to first transpose them:

transposed_blocks = [zip(*row) for row in blocks]

[
    [
        (('.', '4', '.'), ('.', '.', '4'), ('.', '.', '.')),
        (('.', '.', '.'), ('.', '.', '.'), ('1', '.', '.')),
        (('.', '.', '.'), ('.', '.', '.'), ('7', '.', '.'))
    ], [
        (('.', '.', '.'), ('.', '.', '.'), ('.', '.', '.')),
        (('.', '.', '.'), ('3', '.', '.'), ('.', '.', '6')),
        (('.', '.', '.'), ('.', '6', '.'), ('.', '9', '.'))
    ], [
        (('.', '.', '.'), ('.', '.', '.'), ('.', '.', '.')),
        (('.', '1', '.'), ('.', '.', '.'), ('8', '.', '.')),
        (('.', '.', '.'), ('2', '.', '.'), ('.', '.', '.'))
    ]
]

You see how each block is now on its own line. The last step is to merge each line into one list:

blocks_as_list = [reduce(lambda a, b: a+b, item)
                  for l in transposed_blocks for item in l]

[
    ('.', '4', '.', '.', '.', '4', '.', '.', '.'),
    ('.', '.', '.', '.', '.', '.', '1', '.', '.'),
    ('.', '.', '.', '.', '.', '.', '7', '.', '.'),
    ('.', '.', '.', '.', '.', '.', '.', '.', '.'),
    ('.', '.', '.', '3', '.', '.', '.', '.', '6'),
    ('.', '.', '.', '.', '6', '.', '.', '9', '.'),
    ('.', '.', '.', '.', '.', '.', '.', '.', '.'),
    ('.', '1', '.', '.', '.', '.', '8', '.', '.'),
    ('.', '.', '.', '2', '.', '.', '.', '.', '.')
]

And we end up with the list of all blocks. Now you can check them against the sudoku rules.

Is it possible to use "zip()" with a step parameter?

3 Answers3